def test_anon_tuple_with_names_to_container_with_names(self): anon_tuple = structure.Struct([('a', 1), ('b', 2.0)]) types = [('a', tf.int32), ('b', tf.float32)] self.assertDictEqual( type_conversions.type_to_py_container( anon_tuple, computation_types.StructWithPythonType(types, dict)), { 'a': 1, 'b': 2.0 }) self.assertSequenceEqual( type_conversions.type_to_py_container( anon_tuple, computation_types.StructWithPythonType( types, collections.OrderedDict)), collections.OrderedDict([('a', 1), ('b', 2.0)])) test_named_tuple = collections.namedtuple('TestNamedTuple', ['a', 'b']) self.assertSequenceEqual( type_conversions.type_to_py_container( anon_tuple, computation_types.StructWithPythonType(types, test_named_tuple)), test_named_tuple(a=1, b=2.0)) @attr.s class TestFoo(object): a = attr.ib() b = attr.ib() self.assertEqual( type_conversions.type_to_py_container( anon_tuple, computation_types.StructWithPythonType(types, TestFoo)), TestFoo(a=1, b=2.0))
def test_anon_tuple_without_names_promoted_to_container_with_names(self): anon_tuple = structure.Struct([(None, 1), (None, 2.0)]) types = [('a', tf.int32), ('b', tf.float32)] dict_converted_value = type_conversions.type_to_py_container( anon_tuple, computation_types.StructWithPythonType(types, dict)) odict_converted_value = type_conversions.type_to_py_container( anon_tuple, computation_types.StructWithPythonType(types, collections.OrderedDict)) test_named_tuple = collections.namedtuple('TestNamedTuple', ['a', 'b']) named_tuple_converted_value = type_conversions.type_to_py_container( anon_tuple, computation_types.StructWithPythonType(types, test_named_tuple)) @attr.s class TestFoo(object): a = attr.ib() b = attr.ib() attr_converted_value = type_conversions.type_to_py_container( anon_tuple, computation_types.StructWithPythonType(types, TestFoo)) self.assertIsInstance(dict_converted_value, dict) self.assertIsInstance(odict_converted_value, collections.OrderedDict) self.assertIsInstance(named_tuple_converted_value, test_named_tuple) self.assertIsInstance(attr_converted_value, TestFoo)
def test_with_two_level_tuple(self): type_signature = computation_types.StructWithPythonType([ ('a', tf.bool), ('b', computation_types.StructWithPythonType([ ('c', computation_types.TensorType(tf.float32)), ('d', computation_types.TensorType(tf.int32, [20])), ], collections.OrderedDict)), ('e', computation_types.StructType([])), ], collections.OrderedDict) dtypes, shapes = type_conversions.type_to_tf_dtypes_and_shapes( type_signature) self.assert_nested_struct_eq(dtypes, { 'a': tf.bool, 'b': { 'c': tf.float32, 'd': tf.int32 }, 'e': (), }) self.assert_nested_struct_eq( shapes, { 'a': tf.TensorShape([]), 'b': { 'c': tf.TensorShape([]), 'd': tf.TensorShape([20]) }, 'e': (), })
class IsSumCompatibleTest(parameterized.TestCase): @parameterized.named_parameters([ ('tensor_type', computation_types.TensorType(tf.int32)), ('tuple_type_int', computation_types.StructType([tf.int32, tf.int32], )), ('tuple_type_float', computation_types.StructType([tf.complex128, tf.float32, tf.float64])), ('federated_type', computation_types.FederatedType(tf.int32, placements.CLIENTS)), ]) def test_positive_examples(self, type_spec): type_analysis.check_is_sum_compatible(type_spec) @parameterized.named_parameters([ ('tensor_type_bool', computation_types.TensorType(tf.bool)), ('tensor_type_string', computation_types.TensorType(tf.string)), ('partially_defined_shape', computation_types.TensorType(tf.int32, shape=[None])), ('tuple_type', computation_types.StructType([tf.int32, tf.bool])), ('sequence_type', computation_types.SequenceType(tf.int32)), ('placement_type', computation_types.PlacementType()), ('function_type', computation_types.FunctionType(tf.int32, tf.int32)), ('abstract_type', computation_types.AbstractType('T')), ('ragged_tensor', computation_types.StructWithPythonType([], tf.RaggedTensor)), ('sparse_tensor', computation_types.StructWithPythonType([], tf.SparseTensor)), ]) def test_negative_examples(self, type_spec): with self.assertRaises(type_analysis.SumIncompatibleError): type_analysis.check_is_sum_compatible(type_spec)
def test_nested_py_containers(self): anon_tuple = structure.Struct([ (None, 1), (None, 2.0), ('dict_key', structure.Struct([('a', 3), ('b', structure.Struct([(None, 4), (None, 5)]))])) ]) dict_subtype = computation_types.StructWithPythonType( [('a', tf.int32), ('b', computation_types.StructWithPythonType([tf.int32, tf.int32], tuple))], dict) type_spec = computation_types.StructType([(None, tf.int32), (None, tf.float32), ('dict_key', dict_subtype)]) expected_nested_structure = structure.Struct([ (None, 1), (None, 2.0), ('dict_key', { 'a': 3, 'b': (4, 5) }), ]) self.assertEqual( type_conversions.type_to_py_container(anon_tuple, type_spec), expected_nested_structure)
class FnToBuildingBlockTest(parameterized.TestCase): # pyformat: disable @parameterized.named_parameters(( 'nested_fn_same', lambda f, x: f(f(x)), computation_types.StructType( [('f', computation_types.FunctionType(tf.int32, tf.int32)), ('x', tf.int32)]), '(FEDERATED_foo -> (let fc_FEDERATED_symbol_0=FEDERATED_foo.f(FEDERATED_foo.x),fc_FEDERATED_symbol_1=FEDERATED_foo.f(fc_FEDERATED_symbol_0) in fc_FEDERATED_symbol_1))' ), ('nested_fn_different', lambda f, g, x: f(g(x)), computation_types.StructType( [('f', computation_types.FunctionType(tf.int32, tf.int32)), ('g', computation_types.FunctionType(tf.int32, tf.int32)), ('x', tf.int32)]), '(FEDERATED_foo -> (let fc_FEDERATED_symbol_0=FEDERATED_foo.g(FEDERATED_foo.x),fc_FEDERATED_symbol_1=FEDERATED_foo.f(fc_FEDERATED_symbol_0) in fc_FEDERATED_symbol_1))' ), ('selection', lambda x: (x[1], x[0]), computation_types.StructType([tf.int32, tf.int32]), '(FEDERATED_foo -> <FEDERATED_foo[1],FEDERATED_foo[0]>)'), ('constant', lambda: 'stuff', None, '( -> (let fc_FEDERATED_symbol_0=comp#')) # pyformat: enable def test_returns_result(self, fn, parameter_type, fn_str): parameter_name = 'foo' if parameter_type is not None else None result, _ = _federated_computation_serializer(fn, parameter_name, parameter_type) self.assertStartsWith(str(result), fn_str) # pyformat: disable @parameterized.named_parameters( ('tuple', lambda x: (x[1], x[0]), computation_types.StructType([tf.int32, tf.float32]), computation_types.StructWithPythonType([(None, tf.float32), (None, tf.int32)], tuple)), ('list', lambda x: [x[1], x[0]], computation_types.StructType([tf.int32, tf.float32]), computation_types.StructWithPythonType([(None, tf.float32), (None, tf.int32)], list)), ('odict', lambda x: collections.OrderedDict([('A', x[1]), ('B', x[0])]), computation_types.StructType([tf.int32, tf.float32]), computation_types.StructWithPythonType([('A', tf.float32), ('B', tf.int32)], collections.OrderedDict)), ('namedtuple', lambda x: TestNamedTuple(x=x[1], y=x[0]), computation_types.StructType([tf.int32, tf.float32]), computation_types.StructWithPythonType([('x', tf.float32), ('y', tf.int32)], TestNamedTuple)), ) # pyformat: enable def test_returns_result_with_py_container(self, fn, parameter_type, exepcted_result_type): _, type_signature = _federated_computation_serializer( fn, 'foo', parameter_type) self.assertIs(type(type_signature.result), type(exepcted_result_type)) self.assertIs(type_signature.result.python_container, exepcted_result_type.python_container) self.assertEqual(type_signature.result, exepcted_result_type)
def test_metrics_managers_called_with_training_and_evaluation_time_10( self): training_process = mock.create_autospec( iterative_process.IterativeProcess) training_process.initialize.return_value = 'initialize' training_process.initialize.type_signature.return_value = _test_init_fn.type_signature training_process.next.return_value = ('update', {'metric': 1.0}) training_process.next.type_signature = _test_next_fn.type_signature training_selection_fn = mock.MagicMock() evaluation_fn = mock.MagicMock() evaluation_fn.return_value = {'metric': 1.0} evaluation_fn.type_signature = _test_evaluation_fn.type_signature evaluation_selection_fn = mock.MagicMock() metrics_manager = mock.AsyncMock() with mock.patch('time.time') as mock_time: # Since absl.logging.info uses a call to time.time, we mock it out. mock_time.side_effect = [0.0, 10.0] * 3 with mock.patch('absl.logging.info'): training_loop.run_training_process( training_process=training_process, training_selection_fn=training_selection_fn, total_rounds=1, evaluation_fn=evaluation_fn, evaluation_selection_fn=evaluation_selection_fn, metrics_managers=[metrics_manager]) expected_calls = [] metrics = collections.OrderedDict([ ('evaluation/metric', 1.0), ('evaluation/evaluation_time_in_seconds', 10.0), ]) metrics_type = computation_types.StructWithPythonType([ ('evaluation/metric', tf.float32), ('evaluation/evaluation_time_in_seconds', tf.float32), ], collections.OrderedDict) call = mock.call(metrics, metrics_type, 0) expected_calls.append(call) metrics = collections.OrderedDict([ ('metric', 1.0), ('training_time_in_seconds', 10.0), ('round_number', 1), ('evaluation/metric', 1.0), ('evaluation/evaluation_time_in_seconds', 10.0), ]) metrics_type = computation_types.StructWithPythonType([ ('metric', tf.float32), ('training_time_in_seconds', tf.float32), ('round_number', tf.int32), ('evaluation/metric', tf.float32), ('evaluation/evaluation_time_in_seconds', tf.float32), ], collections.OrderedDict) call = mock.call(metrics, metrics_type, 1) expected_calls.append(call) self.assertEqual(metrics_manager.release.call_args_list, expected_calls)
def _parameter_type( parameters, parameter_types: Tuple[computation_types.Type, ...] ) -> Optional[computation_types.Type]: """Bundle any user-provided parameter types into a single argument type.""" parameter_names = [parameter.name for parameter in parameters] if not parameter_types and not parameters: return None if len(parameter_types) == 1: parameter_type = parameter_types[0] if parameter_type is None and not parameters: return None if len(parameters) == 1: return parameter_type # There is a single parameter type but multiple parameters. if not parameter_type.is_struct() or len(parameter_type) != len( parameters): raise TypeError( f'Function with {len(parameters)} parameters must have a parameter ' f'type with the same number of parameters. Found parameter type ' f'{parameter_type}.') name_list_from_types = structure.name_list(parameter_type) if name_list_from_types: if len(name_list_from_types) != len(parameter_type): raise TypeError( 'Types with both named and unnamed fields cannot be unpacked into ' f'argument lists. Found parameter type {parameter_type}.') if set(name_list_from_types) != set(parameter_names): raise TypeError( 'Function argument names must match field names of parameter type. ' f'Found argument names {parameter_names}, which do not match ' f'{name_list_from_types}, the top-level fields of the parameter ' f'type {parameter_type}.') # The provided parameter type has all named fields which exactly match # the names of the function's parameters. return parameter_type else: # The provided parameter type has no named fields. Apply the names from # the function parameters. parameter_types = ( v for (_, v) in structure.to_elements(parameter_type)) return computation_types.StructWithPythonType( list(zip(parameter_names, parameter_types)), collections.OrderedDict) elif len(parameters) == 1: # If there are multiple provided argument types but the function being # decorated only accepts a single argument, tuple the arguments together. return computation_types.to_type(parameter_types) if len(parameters) != len(parameter_types): raise TypeError( f'Function with {len(parameters)} parameters is ' f'incompatible with provided argument types {parameter_types}.') # The function has `n` parameters and `n` parameter types. # Zip them up into a structure using the names from the function as keys. return computation_types.StructWithPythonType( list(zip(parameter_names, parameter_types)), collections.OrderedDict)
def test_with_ragged_tensor(self): t = type_conversions.infer_type( tf.RaggedTensor.from_row_splits([0, 0, 0, 0], [0, 1, 4])) self.assert_types_identical( t, computation_types.StructWithPythonType([ ('flat_values', computation_types.TensorType(tf.int32, [4])), ('nested_row_splits', computation_types.StructWithPythonType([ (None, computation_types.TensorType(tf.int64, [3])) ], tuple)), ], tf.RaggedTensor))
def test_anon_tuple_without_names_to_container_without_names(self): anon_tuple = structure.Struct([(None, 1), (None, 2.0)]) types = [tf.int32, tf.float32] self.assertSequenceEqual( type_conversions.type_to_py_container( anon_tuple, computation_types.StructWithPythonType(types, list)), [1, 2.0]) self.assertSequenceEqual( type_conversions.type_to_py_container( anon_tuple, computation_types.StructWithPythonType(types, tuple)), (1, 2.0))
def test_returns_model_weights_for_model_callable(self): weights_type = model_utils.weights_type_from_model(TestModel) self.assertEqual( computation_types.StructWithPythonType( [('trainable', computation_types.StructWithPythonType([ computation_types.TensorType(tf.float32, [3]), computation_types.TensorType(tf.float32, [1]), ], list)), ('non_trainable', computation_types.StructWithPythonType([ computation_types.TensorType(tf.int32), ], list))], model_utils.ModelWeights), weights_type)
def test_transforms_unnamed_tuple_type_preserving_tuple_container(self): orig_type = computation_types.StructWithPythonType( [tf.int32, tf.float64], tuple) expected_type = computation_types.StructWithPythonType( [tf.float32, tf.float32], tuple) result_type, mutated = type_transformations.transform_type_postorder( orig_type, _convert_tensor_to_float) noop_type, not_mutated = type_transformations.transform_type_postorder( orig_type, _convert_abstract_type_to_tensor) self.assertEqual(result_type, expected_type) self.assertEqual(noop_type, orig_type) self.assertTrue(mutated) self.assertFalse(not_mutated)
def test_anon_tuple_with_names_to_container_without_names_fails(self): anon_tuple = structure.Struct([(None, 1), ('a', 2.0)]) types = [tf.int32, tf.float32] with self.assertRaisesRegex( ValueError, 'contains a mix of named and unnamed elements'): type_conversions.type_to_py_container( anon_tuple, computation_types.StructWithPythonType(types, tuple)) anon_tuple = structure.Struct([('a', 1), ('b', 2.0)]) with self.assertRaisesRegex(ValueError, 'which does not support names'): type_conversions.type_to_py_container( anon_tuple, computation_types.StructWithPythonType(types, list))
def test_succeeds_with_federated_namedtupletype(self): anon_tuple = structure.Struct([(None, 1), (None, 2.0)]) types = [tf.int32, tf.float32] self.assertSequenceEqual( type_conversions.type_to_py_container( anon_tuple, computation_types.FederatedType( computation_types.StructWithPythonType(types, list), placements.SERVER)), [1, 2.0]) self.assertSequenceEqual( type_conversions.type_to_py_container( anon_tuple, computation_types.FederatedType( computation_types.StructWithPythonType(types, tuple), placements.SERVER)), (1, 2.0))
def test_py_named_tuple(self): py_named_tuple_type = collections.namedtuple('test_tuple', ['a']) t = computation_types.StructWithPythonType([('a', tf.int32)], py_named_tuple_type) self.assertIs(t.python_container, py_named_tuple_type) self.assertEqual( repr(t), 'StructType([(\'a\', TensorType(tf.int32))]) as test_tuple')
def test_ordered_dict(self): t = computation_types.StructWithPythonType([('a', tf.int32)], collections.OrderedDict) self.assertIs(t.python_container, collections.OrderedDict) self.assertEqual( repr(t), 'StructType([(\'a\', TensorType(tf.int32))]) as OrderedDict')
def test_not_anon_tuple_passthrough(self): value = (1, 2.0) result = type_conversions.type_to_py_container( (1, 2.0), computation_types.StructWithPythonType([tf.int32, tf.float32], container_type=list)) self.assertEqual(result, value)
def test_ragged_tensor(self): value = structure.Struct([ ('flat_values', [0, 0, 0, 0]), ('nested_row_splits', [[0, 1, 4]]), ]) value_type = computation_types.StructWithPythonType([ ('flat_values', computation_types.TensorType(tf.int32, [4])), ('nested_row_splits', computation_types.StructWithPythonType([ (None, computation_types.TensorType(tf.int64, [3])) ], tuple)), ], tf.RaggedTensor) result = type_conversions.type_to_py_container(value, value_type) self.assertIsInstance(result, tf.RaggedTensor) self.assertAllEqual(result.flat_values, [0, 0, 0, 0]) self.assertEqual(len(result.nested_row_splits), 1) self.assertAllEqual(result.nested_row_splits[0], [0, 1, 4])
def test_struct_with_container_type(self): x = building_blocks.Reference('foo', tf.int32) y = building_blocks.Reference('bar', tf.bool) z = building_blocks.Struct([x, ('y', y)], tuple) self.assertEqual( z.type_signature, computation_types.StructWithPythonType([tf.int32, ('y', tf.bool)], tuple))
def test_py_attr_class(self): @attr.s class TestFoo(object): a = attr.ib() t = computation_types.StructWithPythonType([('a', tf.int32)], TestFoo) self.assertIs(t.python_container, TestFoo) self.assertEqual( repr(t), 'StructType([(\'a\', TensorType(tf.int32))]) as TestFoo')
def _deserialize_type_spec(serialize_type_variable, python_container=None): """Deserialize a `tff.Type` protocol buffer into a python class instance.""" type_spec = type_serialization.deserialize_type( computation_pb2.Type.FromString( serialize_type_variable.read_value().numpy())) if type_spec.is_struct() and python_container is not None: type_spec = computation_types.StructWithPythonType( structure.iter_elements(type_spec), python_container) return type_conversions.type_to_tf_structure(type_spec)
def test_serialize_deserialize_named_tuple_types_py_container(self): # The Py container is destroyed during ser/de. with_container = computation_types.StructWithPythonType( (tf.int32, tf.bool), tuple) p1 = type_serialization.serialize_type(with_container) without_container = type_serialization.deserialize_type(p1) self.assertNotEqual(with_container, without_container) # Not equal. self.assertIsInstance(without_container, computation_types.StructType) self.assertNotIsInstance(without_container, computation_types.StructWithPythonType) with_container.check_equivalent_to(without_container)
def test_takes_namedtuple_polymorphic(self): MyType = collections.namedtuple('MyType', ['x', 'y']) # pylint: disable=invalid-name @tf.function def foo(t): self.assertIsInstance(t, MyType) return t.x + t.y foo = computation_wrapper_instances.tensorflow_wrapper(foo) concrete_fn = foo.fn_for_argument_type( computation_types.StructWithPythonType([('x', tf.int32), ('y', tf.int32)], MyType)) self.assertEqual(concrete_fn.type_signature.compact_representation(), '(<x=int32,y=int32> -> int32)') concrete_fn = foo.fn_for_argument_type( computation_types.StructWithPythonType([('x', tf.float32), ('y', tf.float32)], MyType)) self.assertEqual(concrete_fn.type_signature.compact_representation(), '(<x=float32,y=float32> -> float32)')
def test_client_placed_tuple(self): value = [ structure.Struct([(None, 1), (None, 2)]), structure.Struct([(None, 3), (None, 4)]) ] type_spec = computation_types.FederatedType( computation_types.StructWithPythonType([(None, tf.int32), (None, tf.int32)], tuple), placements.CLIENTS) self.assertEqual([(1, 2), (3, 4)], type_conversions.type_to_py_container( value, type_spec))
def test_ragged_tensor_spec(self): ragged_tensor = tf.RaggedTensor.from_row_splits([0, 0, 0, 0], [0, 1, 4]) ragged_tensor_spec = tf.RaggedTensorSpec.from_value(ragged_tensor) t = computation_types.to_type(ragged_tensor_spec) self.assertEqual(t.python_container, tf.RaggedTensor) self.assertEqual( t.flat_values, computation_types.TensorType(tf.int32, tf.TensorShape(None))) self.assertEqual( t.nested_row_splits, computation_types.StructWithPythonType( [(None, computation_types.TensorType(tf.int64, [None]))], tuple))
def test_with_sparse_tensor(self): # sparse_tensor = [0, 2, 0, 0, 0] sparse_tensor = tf.SparseTensor(indices=[[1]], values=[2], dense_shape=[5]) t = type_conversions.infer_type(sparse_tensor) self.assert_types_identical( t, computation_types.StructWithPythonType([ ('indices', computation_types.TensorType(tf.int64, [1, 1])), ('values', computation_types.TensorType(tf.int32, [1])), ('dense_shape', computation_types.TensorType(tf.int64, [1])), ], tf.SparseTensor))
def test_with_tensor_triple(self): type_signature = computation_types.StructWithPythonType([ ('a', computation_types.TensorType(tf.int32, [5])), ('b', computation_types.TensorType(tf.bool)), ('c', computation_types.TensorType(tf.float32, [3])), ], collections.OrderedDict) tensor_specs = type_conversions.type_to_tf_tensor_specs(type_signature) self.assert_nested_struct_eq( tensor_specs, { 'a': tf.TensorSpec([5], tf.int32), 'b': tf.TensorSpec([], tf.bool), 'c': tf.TensorSpec([3], tf.float32) })
def transform_to_tff_known_type( type_spec: computation_types.Type ) -> Tuple[computation_types.Type, bool]: """Transforms `StructType` to `StructWithPythonType`.""" if type_spec.is_struct() and not type_spec.is_struct_with_python(): field_is_named = tuple( name is not None for name, _ in structure.iter_elements(type_spec)) has_names = any(field_is_named) is_all_named = all(field_is_named) if is_all_named: return computation_types.StructWithPythonType( elements=structure.iter_elements(type_spec), container_type=collections.OrderedDict), True elif not has_names: return computation_types.StructWithPythonType( elements=structure.iter_elements(type_spec), container_type=tuple), True else: raise TypeError( 'Cannot represent TFF type in TF because it contains ' f'partially named structures. Type: {type_spec}') return type_spec, False
def test_without_names(self): expected_structure = ( tf.TensorSpec(shape=(), dtype=tf.bool), tf.TensorSpec(shape=(), dtype=tf.int32), ) type_spec = computation_types.StructWithPythonType( expected_structure, tuple) tf_structure = type_conversions.type_to_tf_structure(type_spec) with tf.Graph().as_default(): ds = tf.data.experimental.from_variant(tf.compat.v1.placeholder( tf.variant, shape=[]), structure=tf_structure) actual_structure = ds.element_spec self.assertEqual(expected_structure, actual_structure)
def test_finalizers_and_unfinalized_metrics_type_mismatch_raises(self): aggregate_factory = aggregation_factory.SumThenFinalizeFactory() metric_finalizers = collections.OrderedDict(num_examples=tf.function( func=lambda x: x)) local_unfinalized_metrics_type = computation_types.StructWithPythonType( collections.OrderedDict(x=tf.TensorSpec(shape=[None, 2], dtype=tf.float32), y=tf.TensorSpec(shape=[None, 1], dtype=tf.float32)), container_type=collections.OrderedDict) with self.assertRaisesRegex( ValueError, 'The metric names in `metric_finalizers` do not match those'): aggregate_factory.create(metric_finalizers, local_unfinalized_metrics_type)