tff.simulation.compose_dataset_computation_with_computation

Builds a new tff.Computation which constructs datasets on clients.

tff.simulation.compose_dataset_computation_with_computation(
    dataset_computation: tff.Computation,
    computation_body: tff.Computation
) -> tff.Computation

Given a tff.Computation that returns a tf.data.Dataset, and a tff.Computation where exactly one of the arguments is a dataset placed on clients of the same type as returned by the tff.Computation, this function will construct a new tff.Computation that accepts a federated set of values of the same type as the parameter of the dataset_computation, maps dataset_computation over these values, and proceeds with the body of computation_body.

For example, if the type signature of dataset_computation is:

(T -> U*)

and the type signature of computation_body is:

({U*}@CLIENTS -> V)

then the returned computation_body type signature will be:

({T}@CLIENTS -> V)

This functionality is useful in several settings:

  • We may want to push some dataset preprocessing to happen on the clients, as opposed to preprocessing happening on the TFF simultation controller. This may be necessary, e.g., in the case that we want to shuffle client examples.
  • We may want to construct the entire dataset on the clients, given a client id. This may be useful in order to speed up distributed simulations, in order to remove a linear cost incurred in constructing and serializing the datasets on the controller.

dataset_computation An instance of tff.Computation which accepts some parameter and returns an element of tff.SequenceType.
computation_body An instance of tff.Computation that accepts exactly one federated dataset, IE, element of type {B*}@CLIENTS, where B is equivalent to the return type of dataset_computation.

A new tff.Computation satisfying the specification above.

TypeError If the arguments are of the wrong types, their TFF type signatures are incompatible with the specification of this function, or if computation_body declares more than one sequence parameter matching the expected dataset type.