def test_nested_py_containers(self):
anon_tuple = anonymous_tuple.AnonymousTuple([
(None, 1), (None, 2.0),
('dict_key',
anonymous_tuple.AnonymousTuple([
('a', 3),
('b', anonymous_tuple.AnonymousTuple([(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 = anonymous_tuple.AnonymousTuple([
(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)
def test_anon_tuple_with_names_to_container_with_names(self):
anon_tuple = anonymous_tuple.AnonymousTuple([('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_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)
test.assert_nested_struct_eq(dtypes, {
'a': tf.bool,
'b': {
'c': tf.float32,
'd': tf.int32
},
'e': (),
})
test.assert_nested_struct_eq(
shapes, {
'a': tf.TensorShape([]),
'b': {
'c': tf.TensorShape([]),
'd': tf.TensorShape([20])
},
'e': (),
})
def test_anon_tuple_without_names_to_container_with_names_fails(self):
anon_tuple = anonymous_tuple.AnonymousTuple([(None, 1), (None, 2.0)])
types = [('a', tf.int32), ('b', tf.float32)]
with self.assertRaisesRegex(ValueError,
'value.*with unnamed elements'):
type_conversions.type_to_py_container(
anon_tuple,
computation_types.StructWithPythonType(types, dict))
with self.assertRaisesRegex(ValueError,
'value.*with unnamed elements'):
type_conversions.type_to_py_container(
anon_tuple,
computation_types.StructWithPythonType(
types, collections.OrderedDict))
test_named_tuple = collections.namedtuple('TestNamedTuple', ['a', 'b'])
with self.assertRaisesRegex(ValueError,
'value.*with unnamed elements'):
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()
with self.assertRaisesRegex(ValueError,
'value.*with unnamed elements'):
type_conversions.type_to_py_container(
anon_tuple,
computation_types.StructWithPythonType(types, TestFoo))
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 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 _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_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_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_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_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_capture_result_with_ragged_tensor(self):
with tf.Graph().as_default() as graph:
type_spec, binding = tensorflow_utils.capture_result_from_graph(
tf.RaggedTensor.from_row_splits([0, 0, 0, 0], [0, 1, 4]), graph)
del binding
self.assert_types_identical(
type_spec,
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_succeeds_with_federated_namedtupletype(self):
anon_tuple = anonymous_tuple.AnonymousTuple([(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),
placement_literals.SERVER)), [1, 2.0])
self.assertSequenceEqual(
type_conversions.type_to_py_container(
anon_tuple,
computation_types.FederatedType(
computation_types.StructWithPythonType(types, tuple),
placement_literals.SERVER)), (1, 2.0))
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_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_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_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 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_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),
placement_literals.CLIENTS)
self.assertEqual([(1, 2), (3, 4)],
type_conversions.type_to_py_container(value, type_spec))
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_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_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)
tensor_specs = type_conversions.type_to_tf_tensor_specs(type_signature)
self.assert_nested_struct_eq(
tensor_specs, {
'a': tf.TensorSpec([], tf.bool),
'b': {
'c': tf.TensorSpec([], tf.float32),
'd': tf.TensorSpec([20], tf.int32)
},
'e': (),
})
def test_stamp_parameter_in_graph_with_struct_with_python_type(self):
with tf.Graph().as_default() as my_graph:
x = self._checked_stamp_parameter(
'foo',
computation_types.StructWithPythonType([('a', tf.int32),
('b', tf.bool)],
collections.OrderedDict))
self.assertIsInstance(x, structure.Struct)
self.assertTrue(len(x), 2)
self._assert_is_placeholder(x.a, 'foo_a:0', tf.int32, [], my_graph)
self._assert_is_placeholder(x.b, 'foo_b:0', tf.bool, [], my_graph)
def test_capture_result_with_sparse_tensor(self):
with tf.Graph().as_default() as graph:
type_spec, binding = tensorflow_utils.capture_result_from_graph(
tf.SparseTensor(indices=[[1]], values=[2], dense_shape=[5]), graph)
del binding
self.assert_types_identical(
type_spec,
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_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)
test.assert_nested_struct_eq(
tensor_specs, {
'a': tf.TensorSpec([5], tf.int32),
'b': tf.TensorSpec([], tf.bool),
'c': tf.TensorSpec([3], tf.float32)
})
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)
