def testNoDispatchableTypes(self): add_dispatcher = _pywrap_python_api_dispatcher.PythonAPIDispatcher( "tf.math.add", math_ops.add, 2, [0, 1], [], False) self.assertEqual(add_dispatcher.Dispatch(1, 2), NotImplemented) concat_dispatcher = _pywrap_python_api_dispatcher.PythonAPIDispatcher( "tf.concat", array_ops.concat, 2, [1], [0], False) self.assertEqual(concat_dispatcher.Dispatch([1], 0), NotImplemented)
def testDispatchParamOutOfRange(self): with self.assertRaisesRegex(ValueError, "index out of range"): _pywrap_python_api_dispatcher.PythonAPIDispatcher( "some_api", None, 5, [0, 1, 5], [2, 3], True) with self.assertRaisesRegex(ValueError, "index out of range"): _pywrap_python_api_dispatcher.PythonAPIDispatcher( "some_api", None, 5, [0, -3], [2, 3], True) with self.assertRaisesRegex(ValueError, "index out of range"): _pywrap_python_api_dispatcher.PythonAPIDispatcher( "some_api", None, 5, [0, 1], [10, 3], True)
def testMultipleDispatchers(self): dispatcher = dispatch.PythonAPIDispatcher('tf.foo', ['x', 'y', 'name'], (None, )) rt_checker = dispatch.MakeInstanceChecker(ragged_tensor.RaggedTensor) rt_x_checker = dispatch.PySignatureChecker([(0, rt_checker)]) rt_y_checker = dispatch.PySignatureChecker([(1, rt_checker)]) f1 = lambda x, y, name=None: 'f1' f2 = lambda x, y, name=None: 'f2' rt = ragged_factory_ops.constant([[1, 2], [3]]) dispatcher.Register(rt_x_checker, f1) dispatcher.Register(rt_y_checker, f2) self.assertEqual(dispatcher.Dispatch((rt, 5), None), 'f1') self.assertEqual(dispatcher.Dispatch(('foo', rt), None), 'f2') self.assertEqual(dispatcher.Dispatch(('foo', ), {'y': rt}), 'f2') self.assertEqual(dispatcher.Dispatch(('foo', 'bar'), None), NotImplemented) with self.assertRaisesRegex( ValueError, 'Multiple dispatch targets .*' r'match the arguments to tf\.foo'): dispatcher.Dispatch((rt, rt), None)
def testDispatcherReturnsNotImplemented(self): dispatcher = _pywrap_python_api_dispatcher.PythonAPIDispatcher( "tf.math.add", math_ops.add, 2, [0, 1], [], False) x = 5 y = Trace("constant", "disabled") z = Trace("constant", "z") self.assertEqual(dispatcher.Dispatch(x, y), NotImplemented) self.assertEqual(dispatcher.Dispatch(y, x), NotImplemented) self.assertEqual(dispatcher.Dispatch(y, z), NotImplemented) self.assertEqual(dispatcher.Dispatch(z, z), Trace("tf.math.add", z, z))
def testSimpleDispatchWithTrace(self): dispatcher = _pywrap_python_api_dispatcher.PythonAPIDispatcher( "tf.math.add", math_ops.add, 2, [0, 1], [], False) x = 5 y = Trace("constant", "y") z = Trace("constant", "z") Trace.log.clear() self.assertEqual(dispatcher.Dispatch(x, y), Trace("tf.math.add", x, y)) self.assertEqual(dispatcher.Dispatch(y, x), Trace("tf.math.add", y, x)) self.assertEqual(dispatcher.Dispatch(y, z), Trace("tf.math.add", y, z)) self.assertEqual(Trace.log, [ "__tf_dispatch__('Trace', 'tf.math.add')", "__tf_dispatch__('Trace', 'tf.math.add')", "__tf_dispatch__('Trace', 'tf.math.add')" ])
def testSimpleDispatchWithWeightedTensor(self): dispatcher = _pywrap_python_api_dispatcher.PythonAPIDispatcher( "tf.math.add", math_ops.add, 2, [0, 1], [], False) x = 5 y = WeightedTensor([1, 2, 3], 0.6) z = WeightedTensor([10, 20, 30], 0.2) x_plus_y = dispatcher.Dispatch(x, y) y_plus_x = dispatcher.Dispatch(y, x) y_plus_z = dispatcher.Dispatch(y, z) self.assertAllEqual(x_plus_y.tensor, [6, 7, 8]) self.assertAllEqual(y_plus_x.tensor, [6, 7, 8]) self.assertAllEqual(y_plus_z.tensor, [11, 22, 33]) self.assertEqual(x_plus_y.weight, 0.6) self.assertEqual(y_plus_x.weight, 0.6) self.assertEqual(y_plus_z.weight, 0.4)
def testDispatchPrecedenceRightToLeft(self): # We use an API for which dispatch is disabled, so all dispatchers get # called (since this test checks the order of the dispatcher list). dispatcher = _pywrap_python_api_dispatcher.PythonAPIDispatcher( "disabled", None, 5, [4, 0, 1], [2, 3], True) t = Trace("constant", "t") t2_1 = Trace2("constant", "t2_1") t2_2 = Trace2("constant", "t2_2") t2b = Trace2B("constant", "t2b") t3 = Trace3("constant", "t3") t4 = Trace4("constant", "t4") # Three dispatchable types, none of which is a subclass of the other: # * precedence is right_to_left (since we set right_to_left=True in the # PtyonAPIDispatcher constructor). (Note: arguments are scanned # right-to-left, but the elements of list arguments are still scanned # left-to-right.) # * duplicates are removed. Trace.log.clear() result = dispatcher.Dispatch(t2_1, t3, [], [t2_2, t3], t4) self.assertEqual(result, NotImplemented) self.assertEqual(Trace.log, [ "__tf_dispatch__('Trace4', 'disabled')", "__tf_dispatch__('Trace2', 'disabled')", "__tf_dispatch__('Trace3', 'disabled')" ]) # Subtypes are moved before their base types. (Note: moving subtypes occurs # *after* we swap the order to be right-to-left; so the dispatch order here # is not what we'd get by just reversing the final dispatch order if # right_to_left were false.) Trace.log.clear() result = dispatcher.Dispatch(t2_1, t3, [t], [t2_2, t, t3, t4], t2b) self.assertEqual(result, NotImplemented) self.assertEqual(Trace.log, [ "__tf_dispatch__('Trace2B', 'disabled')", "__tf_dispatch__('Trace2', 'disabled')", "__tf_dispatch__('Trace3', 'disabled')", "__tf_dispatch__('Trace4', 'disabled')", "__tf_dispatch__('Trace', 'disabled')" ])
def add_type_based_api_dispatcher(target): """Adds a PythonAPIDispatcher to the given TensorFlow API function.""" if hasattr(target, TYPE_BASED_DISPATCH_ATTR): raise ValueError(f"{target} already has a type-based API dispatcher.") _, unwrapped = tf_decorator.unwrap(target) target_argspec = tf_inspect.getargspec(unwrapped) if target_argspec.varargs or target_argspec.keywords: # @TODO(b/194903203) Add v2 dispatch support for APIs that take varargs # and keywords. Examples of APIs that take varargs and kwargs: meshgrid, # einsum, map_values, map_flat_values. return target setattr( target, TYPE_BASED_DISPATCH_ATTR, _api_dispatcher.PythonAPIDispatcher(unwrapped.__name__, target_argspec.args, target_argspec.defaults)) _TYPE_BASED_DISPATCH_SIGNATURES[target] = collections.defaultdict(list) return target
def testListAndUnionDispatch(self): dispatcher = dispatch.PythonAPIDispatcher('tf.foo', ['x', 'ys', 'name'], (None, )) rt_checker = dispatch.MakeInstanceChecker(ragged_tensor.RaggedTensor) tensor_checker = dispatch.MakeInstanceChecker(ops.Tensor) rt_or_t = dispatch.MakeUnionChecker([rt_checker, tensor_checker]) list_of_rt_or_t = dispatch.MakeListChecker(rt_or_t) f1 = lambda x, ys, name=None: 'f1' dispatcher.Register( dispatch.PySignatureChecker([(0, rt_or_t), (1, list_of_rt_or_t)]), f1) rt = ragged_factory_ops.constant([[1, 2], [3]]) t = constant_op.constant(5) self.assertEqual(dispatcher.Dispatch((rt, [t]), None), 'f1') self.assertEqual(dispatcher.Dispatch((rt, [rt]), None), 'f1') self.assertEqual(dispatcher.Dispatch((t, [rt]), None), 'f1') self.assertEqual(dispatcher.Dispatch((rt, []), None), 'f1') self.assertEqual(dispatcher.Dispatch((t, [t, t, rt, t]), None), 'f1') self.assertEqual(dispatcher.Dispatch((rt, [t], 'my_name'), None), 'f1') self.assertEqual(dispatcher.Dispatch((), {'x': rt, 'ys': [t]}), 'f1') self.assertEqual( dispatcher.Dispatch((), { 'x': rt, 'ys': [t], 'name': 'x' }), 'f1') self.assertEqual(dispatcher.Dispatch((t, [t]), None), NotImplemented) self.assertEqual(dispatcher.Dispatch((t, []), None), NotImplemented) self.assertEqual(dispatcher.Dispatch(('foo', [rt]), None), NotImplemented) self.assertEqual(dispatcher.Dispatch(('foo', 'bar'), None), NotImplemented) self.assertEqual(dispatcher.Dispatch(('foo', 'bar', 'baz'), None), NotImplemented)
def testBasicDispatch(self): dispatcher = dispatch.PythonAPIDispatcher('tf.foo', ['x', 'y', 'name'], (None, )) rt_checker = dispatch.MakeInstanceChecker(ragged_tensor.RaggedTensor) f1 = lambda x, y, name=None: 'f1' dispatcher.Register(dispatch.PySignatureChecker([(0, rt_checker)]), f1) rt = ragged_factory_ops.constant([[1, 2], [3]]) self.assertEqual(dispatcher.Dispatch((rt, 5), None), 'f1') self.assertEqual(dispatcher.Dispatch((rt, 5, 'my_name'), None), 'f1') self.assertEqual(dispatcher.Dispatch((), {'x': rt, 'y': 5}), 'f1') self.assertEqual( dispatcher.Dispatch((), { 'x': rt, 'y': 5, 'name': 'x' }), 'f1') self.assertEqual(dispatcher.Dispatch(('foo', rt), None), NotImplemented) self.assertEqual(dispatcher.Dispatch(('foo', 'bar'), None), NotImplemented) self.assertEqual(dispatcher.Dispatch(('foo', 'bar', 'baz'), None), NotImplemented)
def testDispatchPrecedence(self): # We use an API for which dispatch is disabled, so all dispatchers get # called (since this test checks the order of the dispatcher list). dispatcher = _pywrap_python_api_dispatcher.PythonAPIDispatcher( "disabled", None, 5, [0, 1, 4], [2, 3], False) t = Trace("constant", "t") t2_1 = Trace2("constant", "t2_1") t2_2 = Trace2("constant", "t2_2") t2b = Trace2B("constant", "t2b") t3 = Trace3("constant", "t3") t4 = Trace4("constant", "t4") # Three dispatchable types, none of which is a subclass of the other: # * precedence is left-to-right. # * duplicates are removed. Trace.log.clear() result = dispatcher.Dispatch(t2_1, t3, [], [t2_2, t3], t4) self.assertEqual(result, NotImplemented) self.assertEqual(Trace.log, [ "__tf_dispatch__('Trace2', 'disabled')", "__tf_dispatch__('Trace3', 'disabled')", "__tf_dispatch__('Trace4', 'disabled')" ]) # Subtypes are moved before their base types. Trace.log.clear() result = dispatcher.Dispatch(t2_1, t3, [t], [t2_2, t, t3, t4], t2b) self.assertEqual(result, NotImplemented) self.assertEqual(Trace.log, [ "__tf_dispatch__('Trace2B', 'disabled')", "__tf_dispatch__('Trace2', 'disabled')", "__tf_dispatch__('Trace3', 'disabled')", "__tf_dispatch__('Trace4', 'disabled')", "__tf_dispatch__('Trace', 'disabled')" ])