def test_complex_param_tf_computation(self): # The eager executor (inside the local executor stack) has issue with # tf2_computation and the v1 version of SavedModel, hence the test above is # replicated here. MyType = collections.namedtuple('MyType', ['x', 'd']) # pylint: disable=invalid-name @tf.function def foo(t, odict, unnamed_tuple): self.assertIsInstance(t, MyType) self.assertIsInstance(t.d, dict) self.assertIsInstance(odict, collections.OrderedDict) self.assertIsInstance(unnamed_tuple, tuple) return t.x + t.d['y'] + t.d['z'] + odict['o'] + unnamed_tuple[0] args = [ MyType(1, dict(y=2, z=3)), collections.OrderedDict([('o', 0)]), (0,) ] arg_type = [ MyType(tf.int32, collections.OrderedDict(y=tf.int32, z=tf.int32)), collections.OrderedDict([('o', tf.int32)]), (tf.int32,) ] # Explicit type tf_comp = tff.tf_computation(foo, arg_type) self.assertEqual(tf_comp(*args), 6) # Polymorphic tf_comp = tff.tf_computation(foo) self.assertEqual(tf_comp(*args), 6)
def foo(temperatures, threshold): return tff.federated_sum( tff.federated_map( tff.tf_computation( lambda x, y: tf.cast(tf.greater(x, y), tf.int32), [tf.float32, tf.float32]), [temperatures, tff.federated_broadcast(threshold)]))
def test_tf_comp_first_mode_of_usage_as_non_polymorphic_wrapper(self): # Wrapping a lambda with a parameter. foo = tff.tf_computation(lambda x: x > 10, tf.int32) self.assertEqual(str(foo.type_signature), '(int32 -> bool)') self.assertEqual(foo(9), False) self.assertEqual(foo(11), True) # Wrapping an existing Python function with a parameter. bar = tff.tf_computation(tf.add, (tf.int32, tf.int32)) self.assertEqual(str(bar.type_signature), '(<int32,int32> -> int32)') # Wrapping a no-parameter lambda. baz = tff.tf_computation(lambda: tf.constant(10)) self.assertEqual(str(baz.type_signature), '( -> int32)') self.assertEqual(baz(), 10) # Wrapping a no-parameter Python function. def bak_fn(): return tf.constant(10) bak = tff.tf_computation(bak_fn) self.assertEqual(str(bak.type_signature), '( -> int32)') self.assertEqual(bak(), 10)
def test_tf_comp_third_mode_of_usage_as_polymorphic_callable(self): # Wrapping a lambda. foo = tff.tf_computation(lambda x: x > 0) self.assertEqual(foo(-1), False) self.assertEqual(foo(0), False) self.assertEqual(foo(1), True) # Decorating a Python function. @tff.tf_computation def bar(x, y): return x > y self.assertEqual(bar(0, 1), False) self.assertEqual(bar(1, 0), True) self.assertEqual(bar(0, 0), False)
def test_sequence_reduce(self): add_numbers = tff.tf_computation(tf.add, [tf.int32, tf.int32]) @tff.federated_computation(tff.SequenceType(tf.int32)) def foo1(x): return tff.sequence_reduce(x, 0, add_numbers) self.assertEqual(str(foo1.type_signature), '(int32* -> int32)') @tff.federated_computation( tff.FederatedType(tff.SequenceType(tf.int32), tff.SERVER)) def foo2(x): return tff.sequence_reduce(x, 0, add_numbers) self.assertEqual( str(foo2.type_signature), '(int32*@SERVER -> int32@SERVER)') @tff.federated_computation( tff.FederatedType(tff.SequenceType(tf.int32), tff.CLIENTS)) def foo3(x): return tff.sequence_reduce(x, 0, add_numbers) self.assertEqual( str(foo3.type_signature), '({int32*}@CLIENTS -> {int32}@CLIENTS)')
def _(x): return tff.federated_map( tff.tf_computation(lambda x: x > 10, tf.int32), x)
def foo(x, y): return tff.federated_apply( tff.tf_computation(lambda x, y: x > 10, [tf.int32, tf.int32]), [x, y])
def foo(x): return tff.federated_apply( tff.tf_computation(lambda x: x > 10, tf.int32), x)
def foo(x): plus = tff.tf_computation(tf.add, [tf.int32, tf.int32]) return tff.federated_reduce(x, 0, plus)