tff.backends.mapreduce.consolidate_and_extract_local_processing

Consolidates all the local processing in comp.

tff.backends.mapreduce.consolidate_and_extract_local_processing(
    comp, grappler_config_proto
)

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.

    • GENERIC_* intrinsic calls.

    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 CompiledComputations). 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.

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.

An instance of building_blocks.CompiledComputation that holds the TensorFlow section produced by this extraction step, as described above.