async def _ingest(executor, val, type_spec):
"""A coroutine that handles ingestion.
Args:
executor: An instance of `executor_base.Executor`.
val: The first argument to `context_base.Context.ingest()`.
type_spec: The second argument to `context_base.Context.ingest()`.
Returns:
The result of the ingestion.
Raises:
TypeError: If the `val` is not a value of type `type_spec`.
"""
if isinstance(val, executor_value_base.ExecutorValue):
return val
elif (isinstance(val, structure.Struct) and not type_spec.is_federated()):
type_spec.check_struct()
v_elem = structure.to_elements(val)
t_elem = structure.to_elements(type_spec)
if ([k for k, _ in v_elem] != [k for k, _ in t_elem]):
raise TypeError('Value {} does not match type {}.'.format(
val, type_spec))
ingested = []
for (_, v), (_, t) in zip(v_elem, t_elem):
ingested.append(_ingest(executor, v, t))
ingested = await asyncio.gather(*ingested)
return await executor.create_struct(
structure.Struct(
(name, val) for (name, _), val in zip(t_elem, ingested)))
else:
return await executor.create_value(val, type_spec)
def is_argument_struct(arg) -> bool:
"""Determines if 'arg' is interpretable as an argument struct.
Args:
arg: A value or type to test.
Returns:
True iff 'arg' is either a `Struct` in which all unnamed elements
precede named ones, or a `StructType` with this property, or something
that can be converted into the latter by computation_types.to_type().
Raises:
TypeError: If the argument is neither an `structure.Struct`,
nor a type spec.
"""
if isinstance(arg, structure.Struct):
elements = structure.to_elements(arg)
elif isinstance(arg, typed_object.TypedObject):
return is_argument_struct(arg.type_signature)
else:
arg = computation_types.to_type(arg)
if arg.is_struct():
elements = structure.to_elements(arg)
else:
return False
max_unnamed = -1
min_named = len(elements)
for idx, element in enumerate(elements):
if element[0]:
min_named = min(min_named, idx)
else:
max_unnamed = idx
return max_unnamed < min_named
def is_assignable_from(self, source_type: 'Type') -> bool:
if not isinstance(source_type, StructType):
return False
target_elements = structure.to_elements(self)
source_elements = structure.to_elements(source_type)
return ((len(target_elements) == len(source_elements)) and all(
((source_elements[k][0] in [target_elements[k][0], None]) and
target_elements[k][1].is_assignable_from(source_elements[k][1]))
for k in range(len(target_elements))))
def get_curried(fn):
"""Returns a curried version of function `fn` that takes a parameter tuple.
For functions `fn` of types <T1,T2,....,Tn> -> U, the result is a function
of the form T1 -> (T2 -> (T3 -> .... (Tn -> U) ... )).
Note: No attempt is made at avoiding naming conflicts in cases where `fn`
contains references. The arguments of the curriend function are named `argN`
with `N` starting at 0.
Args:
fn: A value of a functional TFF type.
Returns:
A value that represents the curried form of `fn`.
"""
py_typecheck.check_type(fn, value_base.Value)
py_typecheck.check_type(fn.type_signature, computation_types.FunctionType)
py_typecheck.check_type(fn.type_signature.parameter,
computation_types.StructType)
param_elements = structure.to_elements(fn.type_signature.parameter)
references = []
for idx, (_, elem_type) in enumerate(param_elements):
references.append(
building_blocks.Reference('arg{}'.format(idx), elem_type))
result = building_blocks.Call(value_impl.ValueImpl.get_comp(fn),
building_blocks.Struct(references))
for ref in references[::-1]:
result = building_blocks.Lambda(ref.name, ref.type_signature, result)
return value_impl.ValueImpl(result,
value_impl.ValueImpl.get_context_stack(fn))
async def _zip(self, arg, placement, all_equal):
self._check_arg_is_structure(arg)
py_typecheck.check_type(placement, placement_literals.PlacementLiteral)
self._check_strategy_compatible_with_placement(placement)
children = self._target_executors[placement]
cardinality = len(children)
elements = structure.to_elements(arg.internal_representation)
for _, v in elements:
py_typecheck.check_type(v, list)
if len(v) != cardinality:
raise RuntimeError('Expected {} items, found {}.'.format(
cardinality, len(v)))
new_vals = []
for idx in range(cardinality):
new_vals.append(
structure.Struct([(k, v[idx]) for k, v in elements]))
new_vals = await asyncio.gather(
*[c.create_struct(x) for c, x in zip(children, new_vals)])
return FederatedResolvingStrategyValue(
new_vals,
computation_types.FederatedType(computation_types.StructType(
((k, v.member) if k else v.member
for k, v in structure.iter_elements(arg.type_signature))),
placement,
all_equal=all_equal))
def _to_representative_value(type_spec, elements):
"""Convert to a container to a type understood by TF and TFF."""
if type_spec.is_tensor():
return elements
elif type_spec.is_struct():
field_types = structure.to_elements(type_spec)
is_all_named = all([name is not None for name, _ in field_types])
if is_all_named:
if py_typecheck.is_named_tuple(elements):
values = collections.OrderedDict(
(name, _to_representative_value(field_type, e))
for (name, field_type), e in zip(field_types, elements))
return type(elements)(**values)
else:
values = [(name, _to_representative_value(field_type, elements[name]))
for name, field_type in field_types]
return collections.OrderedDict(values)
else:
return tuple(
_to_representative_value(t, e) for t, e in zip(type_spec, elements))
else:
raise ValueError(
'Coercing a dataset with elements of expected type {!s}, '
'produced a value with incompatible type `{!s}. Value: '
'{!s}'.format(type_spec, type(elements), elements))
def test_multiple_named_and_unnamed(self):
v = [(None, 10), ('foo', 20), ('bar', 30)]
x = structure.Struct(v)
self.assertLen(x, 3)
self.assertEqual(x[0], 10)
self.assertEqual(x[1], 20)
self.assertEqual(x[2], 30)
self.assertRaises(IndexError, lambda _: x[3], None)
self.assertEqual(list(iter(x)), [10, 20, 30])
self.assertEqual(dir(x), ['bar', 'foo'])
self.assertEqual(structure.name_list(x), ['foo', 'bar'])
self.assertEqual(x.foo, 20)
self.assertEqual(x.bar, 30)
self.assertRaises(AttributeError, lambda _: x.baz, None)
self.assertEqual(x, structure.Struct([(None, 10), ('foo', 20),
('bar', 30)]))
self.assertNotEqual(
x, structure.Struct([('foo', 10), ('bar', 20), (None, 30)]))
self.assertEqual(structure.to_elements(x), v)
self.assertEqual(
repr(x), 'Struct([(None, 10), (\'foo\', 20), (\'bar\', 30)])')
self.assertEqual(str(x), '<10,foo=20,bar=30>')
with self.assertRaisesRegex(ValueError, 'unnamed'):
structure.to_odict(x)
with self.assertRaisesRegex(ValueError, 'named and unnamed'):
structure.to_odict_or_tuple(x)
def _lines_for_type(type_spec, formatted):
"""Returns a `list` of strings representing the given `type_spec`.
Args:
type_spec: An instance of a TFF `Type`.
formatted: A boolean indicating if the returned string should be
formatted.
"""
if type_spec.is_abstract():
return [type_spec.label]
elif type_spec.is_federated():
member_lines = _lines_for_type(type_spec.member, formatted)
placement_line = '@{}'.format(type_spec.placement)
if type_spec.all_equal:
return _combine([member_lines, [placement_line]])
else:
return _combine([['{'], member_lines, ['}'], [placement_line]])
elif type_spec.is_function():
if type_spec.parameter is not None:
parameter_lines = _lines_for_type(type_spec.parameter,
formatted)
else:
parameter_lines = ['']
result_lines = _lines_for_type(type_spec.result, formatted)
return _combine([['('], parameter_lines, [' -> '], result_lines,
[')']])
elif type_spec.is_struct():
if len(type_spec) == 0: # pylint: disable=g-explicit-length-test
return ['<>']
elements = structure.to_elements(type_spec)
elements_lines = _lines_for_named_types(elements, formatted)
if formatted:
elements_lines = _indent(elements_lines)
lines = [['<', ''], elements_lines, ['', '>']]
else:
lines = [['<'], elements_lines, ['>']]
return _combine(lines)
elif type_spec.is_placement():
return ['placement']
elif type_spec.is_sequence():
element_lines = _lines_for_type(type_spec.element, formatted)
return _combine([element_lines, ['*']])
elif type_spec.is_tensor():
if type_spec.shape.ndims is None:
return ['{!r}[{}]'.format(type_spec.dtype, None)]
elif type_spec.shape.ndims > 0:
def _value_string(value):
return str(value) if value is not None else '?'
value_strings = [
_value_string(e.value) for e in type_spec.shape.dims
]
values_strings = ','.join(value_strings)
return ['{}[{}]'.format(type_spec.dtype.name, values_strings)]
else:
return [type_spec.dtype.name]
else:
raise NotImplementedError('Unexpected type found: {}.'.format(
type(type_spec)))
def make_dummy_element_for_type_spec(type_spec, none_dim_replacement=0):
"""Creates ndarray of zeros corresponding to `type_spec`.
Returns a list containing this ndarray, whose type is *compatible* with, not
necessarily equal to, `type_spec`. This is due to the fact that some
dimensions of `type_spec` may be indeterminate, representing compatibility
of `type_spec` with any number (e.g. leaving a batch dimension indeterminate
to signify compatibility with batches of any size). However a concrete
structure (like the ndarray) must have specified sizes for its dimensions.
So we construct a dummy element where any `None` dimensions of the shape
of `type_spec` are replaced with the value `none_dim_replacement`.The
default value of 0 therefore returns a dummy element of minimal size which
matches `type_spec`.
Args:
type_spec: Instance of `computation_types.Type`, or something convertible to
one by `computation_types.to_type`.
none_dim_replacement: `int` with which to replace any unspecified tensor
dimensions.
Returns:
Returns possibly nested `numpy ndarray`s containing all zeros: a single
`ndarray` if `type_spec` is a `computation_types.TensorType` and a list
of such arrays if `type_spec` is `computation_types.StructType`.
This data structure is of the minimal size necessary in order to be
compatible with `type_spec`.
"""
type_spec = computation_types.to_type(type_spec)
if not type_analysis.contains_only(
type_spec, lambda t: t.is_struct() or t.is_tensor()):
raise ValueError(
'Cannot construct array for TFF type containing anything '
'other than `computation_types.TensorType` or '
'`computation_types.StructType`; you have passed the '
'type {}'.format(type_spec))
py_typecheck.check_type(none_dim_replacement, int)
if none_dim_replacement < 0:
raise ValueError('Please pass nonnegative integer argument as '
'`none_dim_replacement`.')
def _handle_none_dimension(x):
if x is None or (isinstance(x, tf.compat.v1.Dimension)
and x.value is None):
return none_dim_replacement
return x
if type_spec.is_tensor():
dummy_shape = [_handle_none_dimension(x) for x in type_spec.shape]
if type_spec.dtype == tf.string:
return np.empty(dummy_shape, dtype=str)
return np.zeros(dummy_shape, type_spec.dtype.as_numpy_dtype)
elif type_spec.is_struct():
elements = structure.to_elements(type_spec)
elem_list = []
for _, elem_type in elements:
elem_list.append(make_dummy_element_for_type_spec(elem_type))
return elem_list
async def _ingest(executor, val, type_spec):
"""A coroutine that handles ingestion.
Args:
executor: An instance of `executor_base.Executor`.
val: The first argument to `context_base.Context.ingest()`.
type_spec: The second argument to `context_base.Context.ingest()`.
Returns:
The result of the ingestion.
Raises:
TypeError: If the `val` is not a value of type `type_spec`.
"""
if isinstance(val, executor_value_base.ExecutorValue):
return val
elif isinstance(val, ingestable_base.Ingestable):
val_type = val.type_signature
py_typecheck.check_type(val_type, computation_types.Type)
type_spec.check_assignable_from(val_type)
return await val.ingest(executor)
elif (isinstance(val, structure.Struct) and not type_spec.is_federated()):
type_spec.check_struct()
v_elem = structure.to_elements(val)
t_elem = structure.to_elements(type_spec)
if len(v_elem) != len(t_elem):
raise TypeError(
'Value {} does not match type {}: mismatching tuple length.'.
format(val, type_spec))
for ((vk, _), (tk, _)) in zip(v_elem, t_elem):
if vk not in [tk, None]:
raise TypeError(
'Value {} does not match type {}: mismatching tuple element '
'names {} vs. {}.'.format(val, type_spec, vk, tk))
ingested = []
for (_, v), (_, t) in zip(v_elem, t_elem):
ingested.append(_ingest(executor, v, t))
ingested = await asyncio.gather(*ingested)
return await executor.create_struct(
structure.Struct(
(name, val) for (name, _), val in zip(t_elem, ingested)))
else:
return await executor.create_value(val, type_spec)
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 _make(type_spec, next_unused_tensor_index):
if isinstance(type_spec, computation_types.TensorType):
obj = _XlaSerializerTensorArg(type_spec, next_unused_tensor_index)
next_unused_tensor_index = next_unused_tensor_index + 1
return obj, next_unused_tensor_index
py_typecheck.check_type(type_spec, computation_types.StructType)
elements = []
for k, v in structure.to_elements(type_spec):
obj, next_unused_tensor_index = _make(v, next_unused_tensor_index)
elements.append((k, obj))
obj = _XlaSerializerStructArg(type_spec, elements)
return obj, next_unused_tensor_index
def _to_representative_value(type_spec, elements):
"""Convert to a container to a type understood by TF and TFF."""
if type_spec.is_tensor():
return elements
elif type_spec.is_struct_with_python():
if tf.is_tensor(elements):
# In this case we have a singleton tuple tensor that may have been
# unwrapped by tf.data.
elements = [elements]
py_type = computation_types.StructWithPythonType.get_container_type(
type_spec)
field_types = structure.iter_elements(type_spec)
if (issubclass(py_type, collections.abc.Mapping)
or py_typecheck.is_attrs(py_type)):
values = collections.OrderedDict(
(name,
_to_representative_value(field_type, elements[name]))
for name, field_type in field_types)
return py_type(**values)
else:
values = [
_to_representative_value(field_type, e)
for (_, field_type), e in zip(field_types, elements)
]
if py_typecheck.is_named_tuple(py_type):
return py_type(*values)
return py_type(values)
elif type_spec.is_struct():
field_types = structure.to_elements(type_spec)
is_all_named = all([name is not None for name, _ in field_types])
if is_all_named:
if py_typecheck.is_named_tuple(elements):
values = collections.OrderedDict(
(name, _to_representative_value(field_type, e))
for (name,
field_type), e in zip(field_types, elements))
return type(elements)(**values)
else:
values = [
(name,
_to_representative_value(field_type, elements[name]))
for name, field_type in field_types
]
return collections.OrderedDict(values)
else:
return tuple(
_to_representative_value(t, e)
for t, e in zip(type_spec, elements))
else:
raise ValueError(
'Coercing a dataset with elements of expected type {!s}, '
'produced a value with incompatible type `{!s}. Value: '
'{!s}'.format(type_spec, type(elements), elements))
def _to_struct_internal_rep(
*, value: Any, tf_function_cache: MutableMapping[str, Any],
type_spec: computation_types.StructType,
device: tf.config.LogicalDevice) -> structure.Struct:
"""Converts a python container to internal representation for TF executor."""
type_elem = structure.to_elements(type_spec)
value_elem = (structure.to_elements(structure.from_container(value)))
result_elem = []
if len(type_elem) != len(value_elem):
raise TypeError('Expected a {}-element tuple, found {} elements.'.format(
len(type_elem), len(value_elem)))
for (type_name, elem_type), (val_name,
elem_val) in zip(type_elem, value_elem):
if type_name != val_name:
raise TypeError(
'Mismatching element names in type vs. value: {} vs. {}.'.format(
type_name, val_name))
elem_repr = to_representation_for_type(elem_val, tf_function_cache,
elem_type, device)
result_elem.append((type_name, elem_repr))
return structure.Struct(result_elem)
def unpack_args_from_struct(
struct_with_args) -> Tuple[List[Any], Dict[str, Any]]:
"""Extracts argument types from a struct.
Args:
struct_with_args: An instance of either an `struct.Struct` or
computation_types.StructType` (or something convertible to it by
computation_types.to_type()), on which is_argument_struct() is True.
Returns:
A pair (args, kwargs) containing tuple elements from 'struct_with_args'.
Raises:
TypeError: if 'struct_with_args' is of a wrong type.
"""
if not is_argument_struct(struct_with_args):
raise TypeError('Not an argument struct: {}.'.format(struct_with_args))
if isinstance(struct_with_args, structure.Struct):
elements = structure.to_elements(struct_with_args)
elif isinstance(struct_with_args, typed_object.TypedObject):
elements = []
for index, (name, _) in enumerate(
structure.to_elements(struct_with_args.type_signature)):
if name is not None:
elements.append((name, getattr(struct_with_args, name)))
else:
elements.append((None, struct_with_args[index]))
else:
struct_with_args = computation_types.to_type(struct_with_args)
struct_with_args.check_struct()
elements = structure.to_elements(struct_with_args)
args = []
kwargs = {}
for name, value in elements:
if name is not None:
kwargs[name] = value
else:
args.append(value)
return args, kwargs
def replace_empty_leaf_lists_with_numpy_arrays(lists, type_spec):
"""Replaces empty leaf lists in `lists` with numpy arrays.
This function is primarily used to ensure that an appropriate TF dtype is
inferrable for a structure, even if no elements are actually present.
Args:
lists: Output of `make_empty_list_structure_for_element_type_spec`.
type_spec: An instance of `tff.Type` or something convertible to it, as in
`make_empty_list_structure_for_element_type_spec`.
Returns:
The transformed version of `structure`.
Raises:
TypeError: If the `type_spec` is not of a form described above, or if
`lists` is not of a type compatible with `type_spec`.
"""
type_spec = computation_types.to_type(type_spec)
py_typecheck.check_type(type_spec, computation_types.Type)
if type_spec.is_tensor():
py_typecheck.check_type(lists, list)
if len(lists) > 0: # pylint: disable=g-explicit-length-test
return lists
else:
return np.array([], dtype=type_spec.dtype.as_numpy_dtype)
elif type_spec.is_struct():
elements = structure.to_elements(type_spec)
if isinstance(lists, collections.OrderedDict):
to_return = []
for elem_name, elem_type in elements:
elem_val = replace_empty_leaf_lists_with_numpy_arrays(
lists[elem_name], elem_type)
to_return.append((elem_name, elem_val))
return collections.OrderedDict(to_return)
elif isinstance(lists, tuple):
to_return = []
for idx, (_, elem_type) in enumerate(elements):
elem_val = replace_empty_leaf_lists_with_numpy_arrays(
lists[idx], elem_type)
to_return.append(elem_val)
return tuple(to_return)
else:
raise TypeError(
'Invalid nested structure, unexpected container type {}.'.format(
py_typecheck.type_string(type(lists))))
else:
raise TypeError(
'Expected a tensor or struct type, found {}.'.format(type_spec))
def test_empty(self):
v = []
x = structure.Struct(v)
# Explicitly test the implementation of __len__() here so use, assertLen()
# instead of assertEmpty().
self.assertLen(x, 0) # pylint: disable=g-generic-assert
self.assertRaises(IndexError, lambda _: x[0], None)
self.assertEqual(list(iter(x)), [])
self.assertEqual(dir(x), [])
self.assertRaises(AttributeError, lambda _: x.foo, None)
self.assertEqual(x, structure.Struct([]))
self.assertNotEqual(x, structure.Struct([('foo', 10)]))
self.assertEqual(structure.to_elements(x), v)
self.assertEqual(structure.to_odict(x), collections.OrderedDict())
self.assertEqual(repr(x), 'Struct([])')
self.assertEqual(str(x), '<>')
async def compute(self):
if self._type_signature.is_function():
raise TypeError(
'Materializing a computed value of a functional TFF type {} is not '
'possible; only non-functional values can be materialized. Did you '
'perhaps forget to __call__() a function you declared?'.format(
str(self._type_signature)))
elif isinstance(self._value, executor_value_base.ExecutorValue):
return await self._value.compute()
else:
# `ScopedLambda` would have had to declare a functional type, so this is
# the only case left to handle.
py_typecheck.check_type(self._value, structure.Struct)
elem = structure.to_elements(self._value)
vals = await asyncio.gather(*[v.compute() for _, v in elem])
return structure.Struct(zip([k for k, _ in elem], vals))
def _stamp_value_into_graph(value, type_signature, graph):
"""Stamps `value` in `graph` as an object of type `type_signature`.
Args:
value: An value to stamp.
type_signature: An instance of `computation_types.Type`.
graph: The graph to stamp in.
Returns:
A Python object made of tensors stamped into `graph`, `tf.data.Dataset`s,
and `structure.Struct`s that structurally corresponds to the
value passed at input.
"""
py_typecheck.check_type(type_signature, computation_types.Type)
py_typecheck.check_type(graph, tf.Graph)
if value is None:
return None
if type_signature.is_tensor():
if isinstance(value, np.ndarray):
value_type = computation_types.TensorType(
tf.dtypes.as_dtype(value.dtype), tf.TensorShape(value.shape))
type_signature.is_assignable_from(value_type)
with graph.as_default():
return tf.constant(value)
else:
with graph.as_default():
return tf.constant(value,
dtype=type_signature.dtype,
shape=type_signature.shape)
elif type_signature.is_struct():
if isinstance(value, (list, dict)):
value = structure.from_container(value)
stamped_elements = []
named_type_signatures = structure.to_elements(type_signature)
for (name, type_signature), element in zip(named_type_signatures,
value):
stamped_element = _stamp_value_into_graph(element, type_signature,
graph)
stamped_elements.append((name, stamped_element))
return structure.Struct(stamped_elements)
elif type_signature.is_sequence():
return tensorflow_utils.make_data_set_from_elements(
graph, value, type_signature.element)
else:
raise NotImplementedError(
'Unable to stamp a value of type {} in graph.'.format(
type_signature))
def type_to_tf_structure(type_spec: computation_types.Type):
"""Returns nested `tf.data.experimental.Structure` for a given TFF type.
Args:
type_spec: A `computation_types.Type`, the type specification must be
composed of only named tuples and tensors. In all named tuples that appear
in the type spec, all the elements must be named.
Returns:
An instance of `tf.data.experimental.Structure`, possibly nested, that
corresponds to `type_spec`.
Raises:
ValueError: if the `type_spec` is composed of something other than named
tuples and tensors, or if any of the elements in named tuples are unnamed.
"""
py_typecheck.check_type(type_spec, computation_types.Type)
if type_spec.is_tensor():
return tf.TensorSpec(type_spec.shape, type_spec.dtype)
elif type_spec.is_struct():
elements = structure.to_elements(type_spec)
if not elements:
raise ValueError('Empty tuples are unsupported.')
element_outputs = [(k, type_to_tf_structure(v)) for k, v in elements]
named = element_outputs[0][0] is not None
if not all((e[0] is not None) == named for e in element_outputs):
raise ValueError('Tuple elements inconsistently named.')
if type_spec.python_container is None:
if named:
output = collections.OrderedDict(element_outputs)
else:
output = tuple(v for _, v in element_outputs)
else:
container_type = type_spec.python_container
if (py_typecheck.is_named_tuple(container_type)
or py_typecheck.is_attrs(container_type)):
output = container_type(**dict(element_outputs))
elif named:
output = container_type(element_outputs)
else:
output = container_type(e if e[0] is not None else e[1]
for e in element_outputs)
return output
else:
raise ValueError('Unsupported type {}.'.format(
py_typecheck.type_string(type(type_spec))))
def test_single_unnamed(self):
v = [(None, 10)]
x = structure.Struct(v)
self.assertLen(x, 1)
self.assertRaises(IndexError, lambda _: x[1], None)
self.assertEqual(x[0], 10)
self.assertEqual(list(iter(x)), [10])
self.assertEqual(dir(x), [])
self.assertRaises(AttributeError, lambda _: x.foo, None)
self.assertNotEqual(x, structure.Struct([]))
self.assertNotEqual(x, structure.Struct([('foo', 10)]))
self.assertEqual(x, structure.Struct([(None, 10)]))
self.assertNotEqual(x, structure.Struct([(None, 10), ('foo', 20)]))
self.assertEqual(structure.to_elements(x), v)
self.assertEqual(repr(x), 'Struct([(None, 10)])')
self.assertEqual(str(x), '<10>')
with self.assertRaisesRegex(ValueError, 'unnamed'):
structure.to_odict(x)
async def create_value(self, value, type_spec=None):
type_spec = computation_types.to_type(type_spec)
if isinstance(value, computation_impl.ComputationImpl):
return await self.create_value(
computation_impl.ComputationImpl.get_proto(value),
executor_utils.reconcile_value_with_type_spec(value, type_spec))
elif isinstance(value, pb.Computation):
return await self._evaluate(value)
elif type_spec is not None and type_spec.is_struct():
v_el = structure.to_elements(structure.from_container(value))
vals = await asyncio.gather(
*[self.create_value(val, t) for (_, val), t in zip(v_el, type_spec)])
return ReferenceResolvingExecutorValue(
structure.Struct((name, val) for (name, _), val in zip(v_el, vals)))
else:
return ReferenceResolvingExecutorValue(await
self._target_executor.create_value(
value, type_spec))
def make_empty_list_structure_for_element_type_spec(type_spec):
"""Creates a nested structure of empty Python lists for `type_spec`.
This function prepares a nested structure made of `collections.OrderedDict`s
and Python `tuple`s at the intermediate (non-leaf) levels, and that has empty
Python `list`s at the leaf level, with the shape of the structure matching
that of `type_spec`. This structure is used to accumulate elemnts of a data
set for ingestion by `tf.data.Dataset.from_tensor_slices`.
Args:
type_spec: An instance of `tff.Type` or something convertible to it that
consists of only tensor and named tuple types, and in which rach of the
named tuples either have all or none of their elements named.
Returns:
The nested structure, as described above.
Raises:
TypeError: If the `type_spec` is not of a form described above.
"""
type_spec = computation_types.to_type(type_spec)
py_typecheck.check_type(type_spec, computation_types.Type)
if type_spec.is_tensor():
return []
elif type_spec.is_struct():
elements = structure.to_elements(type_spec)
if all(k is not None for k, _ in elements):
return collections.OrderedDict([
(k, make_empty_list_structure_for_element_type_spec(v))
for k, v in elements
])
elif all(k is None for k, _ in elements):
return tuple([
make_empty_list_structure_for_element_type_spec(v)
for _, v in elements
])
else:
raise TypeError(
'Expected a named tuple type with either all elements named or all '
'unnamed, got {}.'.format(type_spec))
else:
raise TypeError(
'Expected a tensor or named tuple type, found {}.'.format(
type_spec))
def test_single_named(self):
v = [('foo', 20)]
x = structure.Struct(v)
self.assertLen(x, 1)
self.assertEqual(x[0], 20)
self.assertRaises(IndexError, lambda _: x[1], None)
self.assertEqual(list(iter(x)), [20])
self.assertEqual(dir(x), ['foo'])
self.assertEqual(x.foo, 20)
self.assertRaises(AttributeError, lambda _: x.bar, None)
self.assertNotEqual(x, structure.Struct([]))
self.assertNotEqual(x, structure.Struct([('foo', 10)]))
self.assertNotEqual(x, structure.Struct([(None, 20)]))
self.assertEqual(x, structure.Struct([('foo', 20)]))
self.assertNotEqual(x, structure.Struct([('foo', 20), ('bar', 30)]))
self.assertEqual(structure.to_elements(x), v)
self.assertEqual(repr(x), 'Struct([(\'foo\', 20)])')
self.assertEqual(str(x), '<foo=20>')
self.assertEqual(structure.to_odict(x), collections.OrderedDict(v))
async def create_struct(self, elements):
if not isinstance(elements, structure.Struct):
elements = structure.from_container(elements)
element_strings = []
element_kv_pairs = structure.to_elements(elements)
to_gather = []
type_elements = []
for k, v in element_kv_pairs:
py_typecheck.check_type(v, CachedValue)
to_gather.append(v.target_future)
if k is not None:
py_typecheck.check_type(k, str)
element_strings.append('{}={}'.format(k, v.identifier))
type_elements.append((k, v.type_signature))
else:
element_strings.append(str(v.identifier))
type_elements.append(v.type_signature)
type_spec = computation_types.StructType(type_elements)
gathered = await asyncio.gather(*to_gather)
identifier = CachedValueIdentifier('<{}>'.format(
','.join(element_strings)))
try:
cached_value = self._cache[identifier]
except KeyError:
target_future = asyncio.ensure_future(
self._target_executor.create_struct(
structure.Struct(
(k, v)
for (k, _), v in zip(element_kv_pairs, gathered))))
cached_value = CachedValue(identifier, None, type_spec,
target_future)
self._cache[identifier] = cached_value
try:
target_value = await cached_value.target_future
except Exception:
# TODO(b/145514490): This is a bit heavy handed, there maybe caches where
# only the current cache item needs to be invalidated; however this
# currently only occurs when an inner RemoteExecutor has the backend go
# down.
self._cache = {}
raise
type_spec.check_assignable_from(target_value.type_signature)
return cached_value
def infer_cardinalities(value, type_spec):
"""Infers cardinalities from Python `value`.
Allows for any Python object to represent a federated value; enforcing
particular representations is not the job of this inference function, but
rather ingestion functions lower in the stack.
Args:
value: Python object from which to infer TFF placement cardinalities.
type_spec: The TFF type spec for `value`, determining the semantics for
inferring cardinalities. That is, we only pull the cardinality off of
federated types.
Returns:
Dict of cardinalities.
Raises:
ValueError: If conflicting cardinalities are inferred from `value`.
TypeError: If the arguments are of the wrong types, or if `type_spec` is
a federated type which is not `all_equal` but the yet-to-be-embedded
`value` is not represented as a Python `list`.
"""
if value is None:
return {}
py_typecheck.check_type(type_spec, computation_types.Type)
if isinstance(value, cardinality_carrying_base.CardinalityCarrying):
return value.cardinality
if type_spec.is_federated():
if type_spec.all_equal:
return {}
if not isinstance(value, (list, tuple)):
raise InvalidNonAllEqualValueError(value, type_spec)
return {type_spec.placement: len(value)}
elif type_spec.is_struct():
structure_value = structure.from_container(value, recursive=False)
cardinality_dict = {}
for idx, (_, elem_type) in enumerate(structure.to_elements(type_spec)):
cardinality_dict = merge_cardinalities(
cardinality_dict, infer_cardinalities(structure_value[idx],
elem_type))
return cardinality_dict
else:
return {}
async def compute_federated_zip_at_clients(
self, arg: FederatedComposingStrategyValue
) -> FederatedComposingStrategyValue:
py_typecheck.check_type(arg.type_signature,
computation_types.StructType)
py_typecheck.check_len(arg.type_signature, 2)
py_typecheck.check_type(arg.internal_representation, structure.Struct)
py_typecheck.check_len(arg.internal_representation, 2)
keys = [k for k, _ in structure.to_elements(arg.type_signature)]
vals = [arg.internal_representation[n] for n in [0, 1]]
types = [arg.type_signature[n] for n in [0, 1]]
for n in [0, 1]:
type_analysis.check_federated_type(
types[n], placement=placement_literals.CLIENTS)
types[n] = computation_types.at_clients(types[n].member)
py_typecheck.check_type(vals[n], list)
py_typecheck.check_len(vals[n], len(self._target_executors))
item_type = computation_types.StructType([
((keys[n], types[n].member) if keys[n] else types[n].member)
for n in [0, 1]
])
result_type = computation_types.at_clients(item_type)
zip_type = computation_types.FunctionType(
computation_types.StructType([
((keys[n], types[n]) if keys[n] else types[n]) for n in [0, 1]
]), result_type)
zip_comp = executor_utils.create_intrinsic_comp(
intrinsic_defs.FEDERATED_ZIP_AT_CLIENTS, zip_type)
async def _child_fn(ex, x, y):
py_typecheck.check_type(x, executor_value_base.ExecutorValue)
py_typecheck.check_type(y, executor_value_base.ExecutorValue)
return await ex.create_call(
await ex.create_value(zip_comp, zip_type), await
ex.create_struct(structure.Struct([(keys[0], x),
(keys[1], y)])))
result = await asyncio.gather(*[
_child_fn(c, x, y)
for c, x, y in zip(self._target_executors, vals[0], vals[1])
])
return FederatedComposingStrategyValue(result, result_type)
async def create_struct(self, elements):
"""Creates a tuple of `elements`.
Args:
elements: As documented in `executor_base.Executor`.
Returns:
An instance of `EagerValue` that represents the constructed tuple.
"""
elements = structure.to_elements(structure.from_container(elements))
val_elements = []
type_elements = []
for k, v in elements:
py_typecheck.check_type(v, EagerValue)
val_elements.append((k, v.internal_representation))
type_elements.append((k, v.type_signature))
return EagerValue(
structure.Struct(val_elements), self._tf_function_cache,
computation_types.StructType([(k, v) if k is not None else v
for k, v in type_elements]))
def _assert_binding_matches_type_and_value(self, binding, type_spec, val,
graph):
"""Asserts that 'bindings' matches the given type, value, and graph."""
self.assertIsInstance(binding, pb.TensorFlow.Binding)
self.assertIsInstance(type_spec, computation_types.Type)
binding_oneof = binding.WhichOneof('binding')
if binding_oneof == 'tensor':
self.assertTrue(tf.is_tensor(val))
if not isinstance(val, tf.Variable):
# We insert a read_value() op for Variables, which produces
# a name we don't control. Otherwise, names should match:
self.assertEqual(binding.tensor.tensor_name, val.name)
self.assertIsInstance(type_spec, computation_types.TensorType)
self.assertEqual(type_spec.dtype, val.dtype.base_dtype)
self.assertEqual(repr(type_spec.shape), repr(val.shape))
elif binding_oneof == 'sequence':
self.assertIsInstance(val,
type_conversions.TF_DATASET_REPRESENTATION_TYPES)
sequence_oneof = binding.sequence.WhichOneof('binding')
self.assertEqual(sequence_oneof, 'variant_tensor_name')
variant_tensor = graph.get_tensor_by_name(
binding.sequence.variant_tensor_name)
op = str(variant_tensor.op.type)
self.assertTrue((op == 'Placeholder') or ('Dataset' in op))
self.assertEqual(variant_tensor.dtype, tf.variant)
self.assertIsInstance(type_spec, computation_types.SequenceType)
self.assertEqual(
computation_types.to_type(val.element_spec), type_spec.element)
elif binding_oneof == 'struct':
self.assertIsInstance(type_spec, computation_types.StructType)
if not isinstance(val, (list, tuple, structure.Struct)):
self.assertIsInstance(val, dict)
if isinstance(val, collections.OrderedDict):
val = list(val.values())
else:
val = [v for _, v in sorted(val.items())]
for idx, e in enumerate(structure.to_elements(type_spec)):
self._assert_binding_matches_type_and_value(binding.struct.element[idx],
e[1], val[idx], graph)
else:
self.fail('Unknown binding.')
def type_to_tf_dtypes_and_shapes(type_spec: computation_types.Type):
"""Returns nested structures of tensor dtypes and shapes for a given TFF type.
The returned dtypes and shapes match those used by `tf.data.Dataset`s to
indicate the type and shape of their elements. They can be used, e.g., as
arguments in constructing an iterator over a string handle.
Args:
type_spec: A `computation_types.Type`, the type specification must be
composed of only named tuples and tensors. In all named tuples that appear
in the type spec, all the elements must be named.
Returns:
A pair of parallel nested structures with the dtypes and shapes of tensors
defined in `type_spec`. The layout of the two structures returned is the
same as the layout of the nested type defined by `type_spec`. Named tuples
are represented as dictionaries.
Raises:
ValueError: if the `type_spec` is composed of something other than named
tuples and tensors, or if any of the elements in named tuples are unnamed.
"""
py_typecheck.check_type(type_spec, computation_types.Type)
if type_spec.is_tensor():
return (type_spec.dtype, type_spec.shape)
elif type_spec.is_struct():
elements = structure.to_elements(type_spec)
if not elements:
output_dtypes = []
output_shapes = []
elif elements[0][0] is not None:
output_dtypes = collections.OrderedDict()
output_shapes = collections.OrderedDict()
for e in elements:
element_name = e[0]
element_spec = e[1]
if element_name is None:
raise ValueError(
'When a sequence appears as a part of a parameter to a section '
'of TensorFlow code, in the type signature of elements of that '
'sequence all named tuples must have their elements explicitly '
'named, and this does not appear to be the case in {}.'
.format(type_spec))
element_output = type_to_tf_dtypes_and_shapes(element_spec)
output_dtypes[element_name] = element_output[0]
output_shapes[element_name] = element_output[1]
else:
output_dtypes = []
output_shapes = []
for e in elements:
element_name = e[0]
element_spec = e[1]
if element_name is not None:
raise ValueError(
'When a sequence appears as a part of a parameter to a section '
'of TensorFlow code, in the type signature of elements of that '
'sequence all named tuples must have their elements explicitly '
'named, and this does not appear to be the case in {}.'
.format(type_spec))
element_output = type_to_tf_dtypes_and_shapes(element_spec)
output_dtypes.append(element_output[0])
output_shapes.append(element_output[1])
if type_spec.python_container is not None:
container_type = type_spec.python_container
def build_py_container(elements):
if (py_typecheck.is_named_tuple(container_type)
or py_typecheck.is_attrs(container_type)):
return container_type(**dict(elements))
else:
return container_type(elements)
output_dtypes = build_py_container(output_dtypes)
output_shapes = build_py_container(output_shapes)
else:
output_dtypes = tuple(output_dtypes)
output_shapes = tuple(output_shapes)
return (output_dtypes, output_shapes)
else:
raise ValueError('Unsupported type {}.'.format(
py_typecheck.type_string(type(type_spec))))
