class IsStructureOfIntegersTest(parameterized.TestCase): @parameterized.named_parameters( ('empty_struct', computation_types.StructType([])), ('int', computation_types.TensorType(tf.int32)), ('ints', computation_types.StructType([tf.int32, tf.int32])), ('nested_struct', computation_types.StructType([ computation_types.TensorType(tf.int32), computation_types.StructType([tf.int32, tf.int32]) ])), ('federated_int_at_clients', computation_types.FederatedType(tf.int32, placements.CLIENTS)), ) def test_returns_true(self, type_spec): self.assertTrue(type_analysis.is_structure_of_integers(type_spec)) @parameterized.named_parameters( ('bool', computation_types.TensorType(tf.bool)), ('float', computation_types.TensorType(tf.float32)), ('string', computation_types.TensorType(tf.string)), ('int_and_bool', computation_types.StructType([tf.int32, tf.bool])), ('nested_struct', computation_types.StructType([ computation_types.TensorType(tf.int32), computation_types.StructType([tf.bool, tf.bool]) ])), ('sequence_of_ints', computation_types.SequenceType(tf.int32)), ('placement', computation_types.PlacementType()), ('function', computation_types.FunctionType(tf.int32, tf.int32)), ('abstract', computation_types.AbstractType('T')), ) def test_returns_false(self, type_spec): self.assertFalse(type_analysis.is_structure_of_integers(type_spec))
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)
class VisitPreorderTest(parameterized.TestCase): # pyformat: disable @parameterized.named_parameters([ ('abstract_type', computation_types.AbstractType('T'), 1), ('nested_function_type', computation_types.FunctionType( computation_types.FunctionType( computation_types.FunctionType(tf.int32, tf.int32), tf.int32), tf.int32), 7), ('named_tuple_type', computation_types.StructType( [tf.int32, tf.bool, computation_types.SequenceType(tf.int32)]), 5), ('placement_type', computation_types.PlacementType(), 1), ]) # pyformat: enable def test_preorder_call_count(self, type_signature, expected_count): class Counter(object): k = 0 def _count_hits(given_type, arg): del given_type # Unused. Counter.k += 1 return arg type_transformations.visit_preorder(type_signature, _count_hits, None) actual_count = Counter.k self.assertEqual(actual_count, expected_count)
def test_transforms_placement_type(self): orig_type = computation_types.PlacementType() expected_type = computation_types.TensorType(tf.float32) result_type, mutated = type_transformations.transform_type_postorder( orig_type, _convert_placement_type_to_tensor) 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_fails_with_bad_types(self): function = computation_types.FunctionType( None, computation_types.TensorType(tf.int32)) federated = computation_types.FederatedType(tf.int32, placements.CLIENTS) tuple_on_function = computation_types.StructType([federated, function]) def foo(x): # pylint: disable=unused-variable del x # Unused. with self.assertRaisesRegex( TypeError, r'you have attempted to create one with the type {int32}@CLIENTS' ): computation_wrapper_instances.tensorflow_wrapper(foo, federated) # pylint: disable=anomalous-backslash-in-string with self.assertRaisesRegex( TypeError, r'you have attempted to create one with the type \( -> int32\)' ): computation_wrapper_instances.tensorflow_wrapper(foo, function) with self.assertRaisesRegex( TypeError, r'you have attempted to create one with the type placement'): computation_wrapper_instances.tensorflow_wrapper( foo, computation_types.PlacementType()) with self.assertRaisesRegex( TypeError, r'you have attempted to create one with the type T'): computation_wrapper_instances.tensorflow_wrapper( foo, computation_types.AbstractType('T')) with self.assertRaisesRegex( TypeError, r'you have attempted to create one with the type <{int32}@CLIENTS,\( ' '-> int32\)>'): computation_wrapper_instances.tensorflow_wrapper( foo, tuple_on_function)
def test_returns_string_for_placement_type(self): type_spec = computation_types.PlacementType() self.assertEqual(type_spec.compact_representation(), 'placement') self.assertEqual(type_spec.formatted_representation(), 'placement')
def test_is_assignable_from(self): t1 = computation_types.PlacementType() t2 = computation_types.PlacementType() self.assertTrue(t1.is_assignable_from(t1)) self.assertTrue(t1.is_assignable_from(t2))
def test_equality(self): t1 = computation_types.PlacementType() t2 = computation_types.PlacementType() self.assertEqual(t1, t2)
def test_construction(self): t1 = computation_types.PlacementType() self.assertEqual(repr(t1), 'PlacementType()') self.assertEqual(str(t1), 'placement')
def create_whimsy_placement_literal(): """Returns a `placements.PlacementLiteral` and type.""" value = placements.SERVER type_signature = computation_types.PlacementType() return value, type_signature
class CreateStructureOfCoroReferencesTest(parameterized.TestCase, unittest.IsolatedAsyncioTestCase): # pyformat: disable @parameterized.named_parameters( ('tensor', _coro(1), computation_types.TensorType(tf.int32), native_platform.CoroValueReference( _coro(1), computation_types.TensorType(tf.int32))), ('federated', _coro(1), computation_types.FederatedType( computation_types.TensorType(tf.int32), placements.SERVER), native_platform.CoroValueReference( _coro(1), computation_types.TensorType(tf.int32))), ('struct_unnamed', _coro((True, 1, 'a')), computation_types.StructType([ (None, computation_types.TensorType(tf.bool)), (None, computation_types.TensorType(tf.int32)), (None, computation_types.TensorType(tf.string)), ]), structure.Struct([ (None, native_platform.CoroValueReference( _coro(True), computation_types.TensorType(tf.bool))), (None, native_platform.CoroValueReference( _coro(1), computation_types.TensorType(tf.int32))), (None, native_platform.CoroValueReference( _coro('a'), computation_types.TensorType(tf.string))), ])), ('struct_named', _coro(collections.OrderedDict([('a', True), ('b', 1), ('c', 'a')])), computation_types.StructType([ ('a', computation_types.TensorType(tf.bool)), ('b', computation_types.TensorType(tf.int32)), ('c', computation_types.TensorType(tf.string)), ]), structure.Struct([ ('a', native_platform.CoroValueReference( _coro(True), computation_types.TensorType(tf.bool))), ('b', native_platform.CoroValueReference( _coro(1), computation_types.TensorType(tf.int32))), ('c', native_platform.CoroValueReference( _coro('a'), computation_types.TensorType(tf.string))), ])), ('struct_nested', _coro( collections.OrderedDict([ ('x', collections.OrderedDict([('a', True), ('b', 1)])), ('y', collections.OrderedDict([('c', 'a')])), ])), computation_types.StructType([ ('x', computation_types.StructType([ ('a', computation_types.TensorType(tf.bool)), ('b', computation_types.TensorType(tf.int32)), ])), ('y', computation_types.StructType([ ('c', computation_types.TensorType(tf.string)), ])), ]), structure.Struct([ ('x', structure.Struct([ ('a', native_platform.CoroValueReference( _coro(True), computation_types.TensorType(tf.bool))), ('b', native_platform.CoroValueReference( _coro(1), computation_types.TensorType(tf.int32))), ])), ('y', structure.Struct([ ('c', native_platform.CoroValueReference( _coro('a'), computation_types.TensorType(tf.string))), ])), ])), ) # pyformat: enable async def test_returns_value(self, coro, type_signature, expected_value): actual_value = native_platform._create_structure_of_coro_references( coro=coro, type_signature=type_signature) if (isinstance(actual_value, structure.Struct) and isinstance(expected_value, structure.Struct)): structure.is_same_structure(actual_value, expected_value) actual_value = structure.flatten(actual_value) expected_value = structure.flatten(expected_value) for a, b in zip(actual_value, expected_value): a = await a.get_value() b = await b.get_value() self.assertEqual(a, b) else: actual_value = await actual_value.get_value() expected_value = await expected_value.get_value() self.assertEqual(actual_value, expected_value) @parameterized.named_parameters( ('none', None), ('bool', True), ('int', 1), ('str', 'a'), ('list', []), ) def test_raises_type_error_with_type_signature(self, type_signature): coro = _coro(1) with self.assertRaises(TypeError): native_platform._create_structure_of_coro_references( coro=coro, type_signature=type_signature) # pyformat: disable @parameterized.named_parameters( ('federated', computation_types.FederatedType( computation_types.TensorType(tf.int32), placements.CLIENTS)), ('function', computation_types.FunctionType(computation_types.TensorType( tf.int32), computation_types.TensorType(tf.int32))), ('placement', computation_types.PlacementType()), ) # pyformat: enable def test_raises_not_implemented_error_with_type_signature( self, type_signature): coro = _coro(1) with self.assertRaises(NotImplementedError): native_platform._create_structure_of_coro_references( coro=coro, type_signature=type_signature) async def test_returned_structure_materialized_sequentially(self): coro = _coro((True, 1, 'a')) type_signature = computation_types.StructType([ (None, computation_types.TensorType(tf.bool)), (None, computation_types.TensorType(tf.int32)), (None, computation_types.TensorType(tf.string)), ]) result = native_platform._create_structure_of_coro_references( coro=coro, type_signature=type_signature) actual_values = [] for value in result: actual_value = await value.get_value() actual_values.append(actual_value) expected_values = [True, 1, 'a'] self.assertEqual(actual_values, expected_values) async def test_returned_structure_materialized_concurrently(self): coro = _coro((True, 1, 'a')) type_signature = computation_types.StructType([ (None, computation_types.TensorType(tf.bool)), (None, computation_types.TensorType(tf.int32)), (None, computation_types.TensorType(tf.string)), ]) result = native_platform._create_structure_of_coro_references( coro=coro, type_signature=type_signature) actual_values = await asyncio.gather(*[v.get_value() for v in result]) expected_values = [True, 1, 'a'] self.assertEqual(actual_values, expected_values)
def test_serialize_type_with_placement(self): actual_proto = type_serialization.serialize_type( computation_types.PlacementType()) expected_proto = pb.Type(placement=pb.PlacementType()) self.assertEqual(actual_proto, expected_proto)
async def create_value( self, value: Any, type_spec: Any = None) -> executor_value_base.ExecutorValue: """Creates an embedded value from the given `value` and `type_spec`. The kinds of supported `value`s are: * An instance of `intrinsic_defs.IntrinsicDef`. * An instance of `placements.PlacementLiteral`. * An instance of `pb.Computation` if of one of the following kinds: intrinsic, lambda, tensorflow, xla, or data. * A Python `list` if `type_spec` is a federated type. Note: The `value` must be a list even if it is of an `all_equal` type or if there is only a single participant associated with the given placement. * A Python value if `type_spec` is a non-functional, non-federated type. Args: value: An object to embed in the executor, one of the supported types defined by above. type_spec: An optional type convertible to instance of `tff.Type` via `tff.to_type`, the type of `value`. Returns: An instance of `executor_value_base.ExecutorValue` representing a value embedded in the `FederatingExecutor` using a particular `FederatingStrategy`. Raises: TypeError: If the `value` and `type_spec` do not match. ValueError: If `value` is not a kind supported by the `FederatingExecutor`. """ type_spec = computation_types.to_type(type_spec) if isinstance(value, intrinsic_defs.IntrinsicDef): type_analysis.check_concrete_instance_of(type_spec, value.type_signature) return self._strategy.ingest_value(value, type_spec) elif isinstance(value, placements.PlacementLiteral): if type_spec is None: type_spec = computation_types.PlacementType() type_spec.check_placement() return self._strategy.ingest_value(value, type_spec) elif isinstance(value, computation_impl.ConcreteComputation): return await self.create_value( computation_impl.ConcreteComputation.get_proto(value), executor_utils.reconcile_value_with_type_spec(value, type_spec)) elif isinstance(value, pb.Computation): deserialized_type = type_serialization.deserialize_type(value.type) if type_spec is None: type_spec = deserialized_type else: type_spec.check_assignable_from(deserialized_type) which_computation = value.WhichOneof('computation') if which_computation in ['lambda', 'tensorflow', 'xla', 'data']: return self._strategy.ingest_value(value, type_spec) elif which_computation == 'intrinsic': if value.intrinsic.uri in FederatingExecutor._FORWARDED_INTRINSICS: return self._strategy.ingest_value(value, type_spec) intrinsic_def = intrinsic_defs.uri_to_intrinsic_def(value.intrinsic.uri) if intrinsic_def is None: raise ValueError('Encountered an unrecognized intrinsic "{}".'.format( value.intrinsic.uri)) return await self.create_value(intrinsic_def, type_spec) else: raise ValueError( 'Unsupported computation building block of type "{}".'.format( which_computation)) elif type_spec is not None and type_spec.is_federated(): return await self._strategy.compute_federated_value(value, type_spec) else: result = await self._unplaced_executor.create_value(value, type_spec) return self._strategy.ingest_value(result, type_spec)
class ContainsOnlyServerPlacedDataTest(parameterized.TestCase): # pyformat: disable @parameterized.named_parameters( ('struct_unnamed', computation_types.StructType([ (None, computation_types.TensorType(tf.bool)), (None, computation_types.TensorType(tf.int32)), (None, computation_types.TensorType(tf.string)), ])), ('struct_named', computation_types.StructType([ ('a', computation_types.TensorType(tf.bool)), ('b', computation_types.TensorType(tf.int32)), ('c', computation_types.TensorType(tf.string)), ])), ('struct_nested', computation_types.StructType([ ('x', computation_types.StructType([ ('a', computation_types.TensorType(tf.bool)), ('b', computation_types.TensorType(tf.int32)), ])), ('y', computation_types.StructType([ ('c', computation_types.TensorType(tf.string)), ])), ])), ('federated_struct', computation_types.FederatedType( computation_types.StructType([ ('a', computation_types.TensorType(tf.bool)), ('b', computation_types.TensorType(tf.int32)), ('c', computation_types.TensorType(tf.string)), ]), placements.SERVER)), ('federated_sequence', computation_types.FederatedType( computation_types.SequenceType( computation_types.TensorType(tf.int32)), placements.SERVER)), ('federated_tensor', computation_types.FederatedType( computation_types.TensorType(tf.int32), placements.SERVER)), ('sequence', computation_types.SequenceType( computation_types.TensorType(tf.int32))), ('tensor', computation_types.TensorType(tf.int32)), ) # pyformat: enable def test_returns_true(self, type_signature): result = federated_context.contains_only_server_placed_data(type_signature) self.assertTrue(result) # pyformat: disable @parameterized.named_parameters( ('federated', computation_types.FederatedType( computation_types.TensorType(tf.int32), placements.CLIENTS)), ('function', computation_types.FunctionType( computation_types.TensorType(tf.int32), computation_types.TensorType(tf.int32))), ('placement', computation_types.PlacementType()), ) # pyformat: enable def test_returns_false(self, type_signature): result = federated_context.contains_only_server_placed_data(type_signature) self.assertFalse(result) @parameterized.named_parameters( ('none', None), ('bool', True), ('int', 1), ('str', 'a'), ('list', []), ) def test_raises_type_error_with_type_signature(self, type_signature): with self.assertRaises(TypeError): federated_context.contains_only_server_placed_data(type_signature)
def deserialize_type( type_proto: Optional[pb.Type]) -> Optional[computation_types.Type]: """Deserializes 'type_proto' as a computation_types.Type. Note: Currently only deserialization for tensor, named tuple, sequence, and function types is implemented. Args: type_proto: An object that supports same interface as `pb.Type` (e.g. pybind backend C++ `Type` protocol buffer messages), or `None`. Returns: The corresponding instance of computation_types.Type (or None if the argument was None). Raises: TypeError: if the argument is of the wrong type. NotImplementedError: for type variants for which deserialization is not implemented. """ if type_proto is None: return None type_variant = type_proto.WhichOneof('type') if type_variant is None: return None elif type_variant == 'tensor': tensor_proto = type_proto.tensor return computation_types.TensorType( dtype=tf.dtypes.as_dtype(tensor_proto.dtype), shape=_to_tensor_shape(tensor_proto)) elif type_variant == 'sequence': return computation_types.SequenceType( deserialize_type(type_proto.sequence.element)) elif type_variant == 'struct': def empty_str_to_none(s): if s == '': # pylint: disable=g-explicit-bool-comparison return None return s return computation_types.StructType( [(empty_str_to_none(e.name), deserialize_type(e.value)) for e in type_proto.struct.element], convert=False) elif type_variant == 'function': return computation_types.FunctionType( parameter=deserialize_type(type_proto.function.parameter), result=deserialize_type(type_proto.function.result)) elif type_variant == 'placement': return computation_types.PlacementType() elif type_variant == 'federated': placement_oneof = type_proto.federated.placement.WhichOneof( 'placement') if placement_oneof == 'value': return computation_types.FederatedType( member=deserialize_type(type_proto.federated.member), placement=placements.uri_to_placement_literal( type_proto.federated.placement.value.uri), all_equal=type_proto.federated.all_equal) else: raise NotImplementedError( 'Deserialization of federated types with placement spec as {} ' 'is not currently implemented yet.'.format(placement_oneof)) else: raise NotImplementedError( 'Unknown type variant {}.'.format(type_variant))
def test_serialize_deserialize_placement_type(self): self._serialize_deserialize_roundtrip_test([ computation_types.PlacementType(), ])
class DatasetDataSourceIteratorTest(parameterized.TestCase, tf.test.TestCase): def test_init_does_not_raise_type_error(self): datasets = [tf.data.Dataset.from_tensor_slices([1, 2, 3])] * 3 federated_type = computation_types.FederatedType( computation_types.SequenceType(tf.int32), placements.CLIENTS) try: native_platform.DatasetDataSourceIterator( datasets=datasets, federated_type=federated_type) except TypeError: self.fail('Raised TypeError unexpectedly.') @parameterized.named_parameters( ('bool', True), ('int', 1), ('str', 'a'), ('list', [True, 1, 'a']), ) def test_init_raises_type_error_with_datasets(self, datasets): federated_type = computation_types.FederatedType( computation_types.SequenceType(tf.int32), placements.CLIENTS) with self.assertRaises(TypeError): native_platform.DatasetDataSourceIterator( datasets=datasets, federated_type=federated_type) # pyformat: disable @parameterized.named_parameters( ('function', computation_types.FunctionType(computation_types.TensorType( tf.int32), computation_types.TensorType(tf.int32))), ('placement', computation_types.PlacementType()), ('sequence', computation_types.SequenceType(tf.int32)), ('struct', computation_types.StructType([ (None, computation_types.TensorType(tf.bool)), (None, computation_types.TensorType(tf.int32)), (None, computation_types.TensorType(tf.string)), ])), ('tensor', computation_types.TensorType(tf.int32)), ) # pyformat: enable def test_init_raises_type_error_with_federated_type(self, federated_type): datasets = [tf.data.Dataset.from_tensor_slices([1, 2, 3])] * 3 with self.assertRaises(TypeError): native_platform.DatasetDataSourceIterator( datasets=datasets, federated_type=federated_type) def test_init_raises_value_error_with_datasets_empty(self): datasets = [] federated_type = computation_types.FederatedType( computation_types.SequenceType(tf.int32), placements.CLIENTS) with self.assertRaises(ValueError): native_platform.DatasetDataSourceIterator( datasets=datasets, federated_type=federated_type) def test_init_raises_value_error_with_datasets_different_types(self): datasets = [ tf.data.Dataset.from_tensor_slices([1, 2, 3]), tf.data.Dataset.from_tensor_slices(['a', 'b', 'c']), ] federated_type = computation_types.FederatedType( computation_types.SequenceType(tf.int32), placements.CLIENTS) with self.assertRaises(ValueError): native_platform.DatasetDataSourceIterator( datasets=datasets, federated_type=federated_type) @parameterized.named_parameters( ('1', 0), ('2', 1), ('3', 2), ) def test_select_returns_data(self, number_of_clients): datasets = [tf.data.Dataset.from_tensor_slices([1, 2, 3])] * 3 federated_type = computation_types.FederatedType( computation_types.SequenceType(tf.int32), placements.CLIENTS) iterator = native_platform.DatasetDataSourceIterator( datasets=datasets, federated_type=federated_type) data = iterator.select(number_of_clients) self.assertLen(data, number_of_clients) for actual_dataset in data: expected_dataset = tf.data.Dataset.from_tensor_slices([1, 2, 3]) self.assertSameElements(actual_dataset, expected_dataset) @parameterized.named_parameters( ('none', None), ('negative', -1), ('length', 4), ) def test_select_raises_value_error(self, number_of_clients): datasets = [tf.data.Dataset.from_tensor_slices([1, 2, 3])] * 3 federated_type = computation_types.FederatedType( computation_types.SequenceType(tf.int32), placements.CLIENTS) iterator = native_platform.DatasetDataSourceIterator( datasets=datasets, federated_type=federated_type) with self.assertRaises(ValueError): iterator.select(number_of_clients)