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())