def test_with_structure_replacing_federated_zip(self):
fed_tuple = building_blocks.Reference(
'tup',
computation_types.FederatedType([tf.int32] * 3,
placements.CLIENTS))
unzipped = building_block_factory.create_federated_unzip(fed_tuple)
zipped = building_block_factory.create_federated_zip(unzipped)
placement_unwrapped, _ = tree_transformations.unwrap_placement(zipped)
placement_gone = placement_unwrapped.argument
lambdas_and_blocks_removed, modified = transformations.remove_lambdas_and_blocks(
placement_gone)
self.assertTrue(modified)
self.assertNoLambdasOrBlocks(lambdas_and_blocks_removed)
def consolidate_and_extract_local_processing(comp, grappler_config_proto):
"""Consolidates all the local processing in `comp`.
The input computation `comp` must have the following properties:
1. The output of `comp` may be of a federated type or unplaced. We refer to
the placement `p` of that type as the placement of `comp`. There is no
placement anywhere in the body of `comp` different than `p`. If `comp`
is of a functional type, and has a parameter, the type of that parameter
is a federated type placed at `p` as well, or unplaced if the result of
the function is unplaced.
2. The only intrinsics that may appear in the body of `comp` are those that
manipulate data locally within the same placement. The exact set of these
intrinsics will be gradually updated. At the moment, we support only the
following:
* Either `federated_apply` or `federated_map`, depending on whether `comp`
is `SERVER`- or `CLIENTS`-placed. `federated_map_all_equal` is also
allowed in the `CLIENTS`-placed case.
* Either `federated_value_at_server` or `federated_value_at_clients`,
likewise placement-dependent.
* Either `federated_zip_at_server` or `federated_zip_at_clients`, again
placement-dependent.
Anything else, including `sequence_*` operators, should have been reduced
already prior to calling this function.
3. There are no lambdas in the body of `comp` except for `comp` itself being
possibly a (top-level) lambda. All other lambdas must have been reduced.
This requirement may eventually be relaxed by embedding lambda reducer into
this helper method.
4. If `comp` is of a functional type, it is either an instance of
`building_blocks.CompiledComputation`, in which case there is nothing for
us to do here, or a `building_blocks.Lambda`.
5. There is at most one unbound reference under `comp`, and this is only
allowed in the case that `comp` is not of a functional type.
Aside from the intrinsics specified above, and the possibility of allowing
lambdas, blocks, and references given the constraints above, the remaining
constructs in `comp` include a combination of tuples, selections, calls, and
sections of TensorFlow (as `CompiledComputation`s). This helper function does
contain the logic to consolidate these constructs.
The output of this transformation is always a single section of TensorFlow,
which we henceforth refer to as `result`, the exact form of which depends on
the placement of `comp` and the presence or absence of an argument.
a. If there is no argument in `comp`, and `comp` is `SERVER`-placed, then
the `result` is such that `comp` can be equivalently represented as:
```
federated_value_at_server(result())
```
b. If there is no argument in `comp`, and `comp` is `CLIENTS`-placed, then
the `result` is such that `comp` can be equivalently represented as:
```
federated_value_at_clients(result())
```
c. If there is an argument in `comp`, and `comp` is `SERVER`-placed, then
the `result` is such that `comp` can be equivalently represented as:
```
(arg -> federated_apply(<result, arg>))
```
d. If there is an argument in `comp`, and `comp` is `CLIENTS`-placed, then
the `result` is such that `comp` can be equivalently represented as:
```
(arg -> federated_map(<result, arg>))
```
If the type of `comp` is `T@p` (thus `comp` is non-functional), the type of
`result` is `T`, where `p` is the specific (concrete) placement of `comp`.
If the type of `comp` is `(T@p -> U@p)`, then the type of `result` must be
`(T -> U)`, where `p` is again a specific placement.
Args:
comp: An instance of `building_blocks.ComputationBuildingBlock` that serves
as the input to this transformation, as described above.
grappler_config_proto: An instance of `tf.compat.v1.ConfigProto` to
configure Grappler graph optimization of the generated TensorFlow graph.
If `grappler_config_proto` has
`graph_options.rewrite_options.disable_meta_optimizer=True`, Grappler is
bypassed.
Returns:
An instance of `building_blocks.CompiledComputation` that holds the
TensorFlow section produced by this extraction step, as described above.
"""
py_typecheck.check_type(comp, building_blocks.ComputationBuildingBlock)
if comp.type_signature.is_function():
if comp.is_compiled_computation():
return comp
elif not comp.is_lambda():
# We normalize on lambdas for ease of calling unwrap_placement below.
# The constructed lambda here simply forwards its argument to `comp`.
arg = building_blocks.Reference(
next(building_block_factory.unique_name_generator(comp)),
comp.type_signature.parameter)
called_fn = building_blocks.Call(comp, arg)
comp = building_blocks.Lambda(arg.name, arg.type_signature,
called_fn)
if comp.type_signature.result.is_federated():
unwrapped, _ = tree_transformations.unwrap_placement(comp.result)
# Unwrapped can be a call to `federated_value_at_P`, or
# `federated_apply/map`.
if unwrapped.function.uri in (intrinsic_defs.FEDERATED_APPLY.uri,
intrinsic_defs.FEDERATED_MAP.uri):
extracted = parse_tff_to_tf(unwrapped.argument[0],
grappler_config_proto)
check_extraction_result(unwrapped.argument[0], extracted)
return extracted
else:
member_type = None if comp.parameter_type is None else comp.parameter_type.member
rebound = building_blocks.Lambda(comp.parameter_name,
member_type,
unwrapped.argument)
extracted = parse_tff_to_tf(rebound, grappler_config_proto)
check_extraction_result(rebound, extracted)
return extracted
else:
extracted = parse_tff_to_tf(comp, grappler_config_proto)
check_extraction_result(comp, extracted)
return extracted
elif comp.type_signature.is_federated():
unwrapped, _ = tree_transformations.unwrap_placement(comp)
# Unwrapped can be a call to `federated_value_at_P`, or
# `federated_apply/map`.
if unwrapped.function.uri in (intrinsic_defs.FEDERATED_APPLY.uri,
intrinsic_defs.FEDERATED_MAP.uri):
extracted = parse_tff_to_tf(unwrapped.argument[0],
grappler_config_proto)
check_extraction_result(unwrapped.argument[0], extracted)
return extracted
else:
extracted = parse_tff_to_tf(unwrapped.argument,
grappler_config_proto)
check_extraction_result(unwrapped.argument, extracted)
return extracted.function
else:
called_tf = parse_tff_to_tf(comp, grappler_config_proto)
check_extraction_result(comp, called_tf)
return called_tf.function