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)
