def test_defaultbackend_autogradcpu(self):
dispatcher = PythonDispatcher()
dispatcher.register(
["CPU", "XLA", "CompositeExplicitAutograd", "AutogradCPU"])
self.assertExpectedInline(
dispatcher.dispatchTable(), '''\
Computed Dispatch Table
key kernel
---------------------------
CPU fn_CPU [kernel]
XLA fn_XLA [kernel]
QuantizedCPU fn_CompositeExplicitAutograd [default backend kernel]
AutogradOther fallthrough [backend fallback]
AutogradCPU fn_AutogradCPU [kernel]
AutogradXLA fallthrough [backend fallback]
''')
self.assertExpectedInline(
dispatcher.registrations(), '''\
Registered Kernels
key kernel
---------------------------
CPU fn_CPU
XLA fn_XLA
AutogradCPU fn_AutogradCPU
CompositeExplicitAutograd[alias] fn_CompositeExplicitAutograd
''')
def test_autogradother(self):
dispatcher = PythonDispatcher()
dispatcher.register(
["CPU", "QuantizedCPU", "CompositeImplicitAutograd"])
self.assertExpectedInline(
dispatcher.dispatchTable(), '''\
Computed Dispatch Table
key kernel
---------------------------
CPU fn_CPU [kernel]
XLA fn_CompositeImplicitAutograd [math kernel]
QuantizedCPU fn_QuantizedCPU [kernel]
AutogradOther ambiguous_autogradother [ambiguous autogradother]
AutogradCPU fallthrough [backend fallback]
AutogradXLA fn_CompositeImplicitAutograd [math kernel]
''')
self.assertExpectedInline(
dispatcher.registrations(), '''\
Registered Kernels
key kernel
---------------------------
CPU fn_CPU
QuantizedCPU fn_QuantizedCPU
CompositeImplicitAutograd[alias] fn_CompositeImplicitAutograd
''')
def test_math_autogradcpu(self):
dispatcher = PythonDispatcher()
dispatcher.register(
["CPU", "XLA", "Lazy", "CompositeImplicitAutograd", "AutogradCPU"])
self.assertExpectedInline(
dispatcher.dispatchTable(), '''\
Computed Dispatch Table
key kernel
---------------------------
CPU fn_CPU [kernel]
XLA fn_XLA [kernel]
Lazy fn_Lazy [kernel]
FPGA fn_CompositeImplicitAutograd [math kernel]
AutogradOther fn_CompositeImplicitAutograd [math kernel]
AutogradCPU fn_AutogradCPU [kernel]
AutogradXLA fallthrough [backend fallback]
AutogradLazy fallthrough [backend fallback]
''')
self.assertExpectedInline(
dispatcher.registrations(), '''\
Registered Kernels
key kernel
---------------------------
CPU fn_CPU
XLA fn_XLA
Lazy fn_Lazy
AutogradCPU fn_AutogradCPU
CompositeImplicitAutograd[alias] fn_CompositeImplicitAutograd
''')
def test_math_autogradcpu(self):
dispatcher = PythonDispatcher()
dispatcher.register(["CPU", "XLA", "Math", "AutogradCPU"])
self.assertExpectedInline(
dispatcher.dispatchTable(), '''\
Computed Dispatch Table
key kernel
---------------------------
CPU fn_CPU [kernel]
XLA fn_XLA [kernel]
QuantizedCPU fn_Math [math kernel]
AutogradOther fn_Math [math kernel]
AutogradCPU fn_AutogradCPU [kernel]
AutogradXLA fallthrough [backend fallback]
''')
self.assertExpectedInline(
dispatcher.registrations(), '''\
Registered Kernels
key kernel
---------------------------
CPU fn_CPU
XLA fn_XLA
AutogradCPU fn_AutogradCPU
Math[alias] fn_Math
''')
def test_defaultbackend_math(self):
dispatcher = PythonDispatcher()
with self.assertRaisesRegex(
RuntimeError,
r"Registration to both CompositeImplicitAutograd and CompositeExplicitAutograd is not allowed"):
dispatcher.register(["CompositeExplicitAutograd", "CompositeImplicitAutograd"])
def test_defaultbackend_math(self):
dispatcher = PythonDispatcher()
with self.assertRaisesRegex(
RuntimeError,
r"Registration to both Math and DefaultBackend is not allowed"
):
dispatcher.register(["DefaultBackend", "Math"])
def test_basic(self):
dispatcher = PythonDispatcher()
dispatcher.register(["CPU", "XLA", "CompositeImplicitAutograd"])
self.assertExpectedInline(
dispatcher.dispatchTable(), '''\
Computed Dispatch Table
key kernel
---------------------------
CPU fn_CPU [kernel]
XLA fn_XLA [kernel]
QuantizedCPU fn_CompositeImplicitAutograd [math kernel]
AutogradOther fn_CompositeImplicitAutograd [math kernel]
AutogradCPU fallthrough [backend fallback]
AutogradXLA fallthrough [backend fallback]
''')
def test_duplicate_registrations(self):
dispatcher = PythonDispatcher()
with self.assertRaisesRegex(RuntimeError, r"Overriden is not allowed"):
dispatcher.register(["CPU", "CPU"])
