def contains_only_server_placed_data(
type_signature: computation_types.Type) -> bool:
"""Determines if `type_signature` contains only server-placed data.
Determines if `type_signature` contains only:
* `tff.StructType`s
* `tff.SequenceType`s
* server-placed `tff.FederatedType`s
* `tff.TensorType`s
Args:
type_signature: The `tff.Type` to test.
Returns:
`True` if `type_signature` contains only server-placed data, otherwise
`False`.
"""
py_typecheck.check_type(type_signature, computation_types.Type)
def predicate(type_spec: computation_types.Type) -> bool:
return (type_spec.is_struct() or
(type_spec.is_federated() and
type_spec.placement is placements.SERVER) or
type_spec.is_sequence() or type_spec.is_tensor())
return type_analysis.contains_only(type_signature, predicate)
async def _zip_struct_into_child(self, child, child_index, value,
value_type):
"""Embeds `value` elements at `child_index` into `child`."""
if value_type.is_federated():
py_typecheck.check_type(value, list)
return value[child_index]
value_type.check_struct()
# Note that structs whose members are all federated values may appear
# embedded iff the struct *has no non-struct members*, e.g.:
# <>; <<>>; <<>, <>>
if isinstance(value, executor_value_base.ExecutorValue):
if not type_analysis.contains_only(
value_type, lambda element_type: element_type.is_struct()):
raise ValueError(
f'Expected zippable value of type `structure.Struct`, but found '
f'value of type {value} (TFF type {value_type})')
return value
py_typecheck.check_type(value, structure.Struct)
new_elements = await asyncio.gather(*[
self._zip_struct_into_child(child, child_index, element,
element_type)
for element, element_type in zip(value, value_type)
])
named_elements = structure.Struct(
list(zip(structure.name_list_with_nones(value), new_elements)))
return await child.create_struct(named_elements)
def is_stateful_process(process: measured_process.MeasuredProcess) -> bool:
"""Determine if a `MeasuredProcess` has a non-empty state."""
def federated_empty_struct(type_spec: computation_types.Type) -> bool:
return type_spec.is_struct() or type_spec.is_federated()
# Check if any child type is not an empty struct.
return not type_analysis.contains_only(
process.initialize.type_signature.result, federated_empty_struct)
def _build_reservoir_type(
sample_value_type: computation_types.Type) -> computation_types.Type:
"""Create the TFF type for the reservoir's state.
`UnweightedReservoirSamplingFactory` will use this type as the "state" type in
a `tff.federated_aggregate` (an input to `accumulate`, `merge` and `report`).
Args:
sample_value_type: The `tff.Type` of the values that will be aggregated from
clients.
Returns:
A `collection.OrderedDict` with three keys:
random_seed: A 2-tuple of `tf.int64` scalars. This keeps track of the
`seed` parameter for `tf.random.stateless_uniform` calls for
sampling during aggregation.
random_values: A 1-d tensor of `int32` values randomly generated from
`tf.random.stateless_uniform`. The size of the tensor will be the
same as the first dimenion of each leaf in `samples` (these can be
thought of as parallel lists). These values are used to determine
whether a sample stays in the reservoir, or is evicted, as the values
are aggregated. If the i-th value of this list is evicted, then the
i-th (in the first dimension) tensor in each of the `samples` structure
leaves must be evicted.
samples: A tensor or structure of tensors representing the actual sampled
values. If a structure, the shape of the structure matches that of
`sample_value_type`. All tensors have one additional dimenion prepended
which has an unknown size. This will be used to concatenate samples and
store them in the reservoir.
"""
# TODO(b/181365504): relax this to allow `StructType` once a `Struct` can be
# returned from `tf.function` decorated methods.
def is_tesnor_or_struct_with_py_type(t: computation_types.Type) -> bool:
return t.is_tensor() or t.is_struct_with_python()
if not type_analysis.contains_only(sample_value_type,
is_tesnor_or_struct_with_py_type):
raise TypeError(
'Cannot create a reservoir for type structure. Sample type '
'must only contain `TensorType` or `StructWithPythonType`, '
f'got a {sample_value_type!r}.')
def add_unknown_dimension(t):
if t.is_tensor():
return (computation_types.TensorType(dtype=t.dtype,
shape=[None] + t.shape), True)
return t, False
# TODO(b/181155367): creating a value from a type for the `zero` is a common
# pattern for users of `tff.federated_aggregate` that could be made easier
# for TFF users. Replace this once such helper exists.
return computation_types.to_type(
collections.OrderedDict(
random_seed=computation_types.TensorType(tf.int64, shape=[2]),
random_values=computation_types.TensorType(tf.int32, shape=[None]),
samples=type_transformations.transform_type_postorder(
sample_value_type, add_unknown_dimension)[0]))
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
def _is_tensor_or_structure_of_tensors(
value_type: computation_types.Type) -> bool:
"""Return True if `value_type` is a TensorType or structure of TensorTypes."""
# TODO(b/181365504): relax this to allow `StructType` once a `Struct` can be
# returned from `tf.function` decorated methods.
def is_tensor_or_struct_with_py_type(t: computation_types.Type) -> bool:
return t.is_tensor() or t.is_struct_with_python()
return type_analysis.contains_only(value_type,
is_tensor_or_struct_with_py_type)
async def _invoke_mergeable_comp_form(
comp: MergeableCompForm, arg: Optional[MergeableCompExecutionContextValue],
execution_contexts: Sequence[
async_execution_context.AsyncExecutionContext]):
"""Invokes `comp` on `arg`, repackaging the results to a single value."""
if arg is not None:
arg_list = arg.value()
else:
arg_list = [None for _ in range(len(execution_contexts))]
up_to_merge_futures = asyncio.as_completed([
_invoke_up_to_merge_and_return_context(comp, x, context)
for x, context in zip(arg_list, execution_contexts)
])
# We compute merge using the most-recently completed context.
# TODO(b/195349085): merge in a hierarchical fashion here rather than
# linearly.
merge_result, merge_context = await next(up_to_merge_futures)
for up_to_merge_result_future in up_to_merge_futures:
to_merge, merge_context = await up_to_merge_result_future
merge_result = await _merge_results(comp, merge_result, to_merge,
merge_context)
if type_analysis.contains_only(
comp.after_merge.type_signature.result,
lambda x: not x.is_federated() or x.all_equal):
# In this case, all contexts must return the same result, which must
# therefore be independent of which element in the arg_list is passed--so we
# avoid the extra work.
top_level_arg_slice = arg_list[0]
return await _compute_after_merged(comp, top_level_arg_slice,
merge_result, merge_context)
after_merge_results = await asyncio.gather(*[
_compute_after_merged(comp, x, merge_result, context)
for x, context in zip(arg_list, execution_contexts)
])
repackaged_values = _repackage_partitioned_values(
after_merge_results,
result_type_spec=comp.after_merge.type_signature.result)
return repackaged_values
def is_stateful(process: IterativeProcess) -> bool:
"""Determines whether a process is stateful.
A process that has a non-empty state is called "stateful"; it follows that
process with an empty state is called "stateless". In TensorFlow Federated
"empty" means a type structure that contains only `tff.types.StructType`; no
tensors or other values. These structures are "empty" in the sense no values
need be communicated and flattening the structure would result in an empty
list.
Args:
process: The `IterativeProcess` to test for statefulness.
Returns:
`True` iff the process is stateful and has a state type structure that
contains types other than `tff.types.StructType`, `False` otherwise.
"""
state_type = process.state_type
if state_type.is_federated():
state_type = state_type.member
return not type_analysis.contains_only(state_type, lambda t: t.is_struct())
def test_returns_false(self, type_signature, types):
result = type_analysis.contains_only(type_signature,
lambda x: isinstance(x, types))
self.assertFalse(result)
def _only_tuple_or_tensor(x: computation_types.Type) -> bool:
return type_analysis.contains_only(
x, lambda t: t.is_struct() or t.is_tensor())
def _only_tuple_or_tensor(value):
return type_analysis.contains_only(
value.type_signature, lambda t: t.is_struct() or t.is_tensor())
def _is_compatible(t: computation_types.Type) -> bool:
return type_analysis.contains_only(
t,
lambda t: t.is_struct() or t.is_tensor() or t.is_federated())
