tff.framework.MergeableCompForm

A data class for computations containing a single logical aggregation.

tff.framework.MergeableCompForm(
    *,
    up_to_merge: tff.Computation,
    merge: tff.Computation,
    after_merge: tff.Computation
)

MergeableCompForm contains three member computations, up_to_merge and merge, and after_merge. A computation in MergeableCompForm is defined to be equivalent to invoking up_to_merge on subsets of CLIENTS-placed arguments in sequence, invoking merge on the stream of these results, and passing the resulting value (placed at tff.SERVER) to after_merge, in addition to the original argument to up_to_merge. In the case of a no-arg up_to_merge computation, after_merge should accept only the result of merge. up_to_merge should return a single tff.SERVER-placed value.

We require that computations in MergeableCompForm contain no AST nodes with signatures taking arguments at tff.CLIENTS and producing results at tff.SERVER.

MergeableCompForm computations are often generated by aligning a computation containing a single logical aggregation on this aggregation, and splitting it along its merge boundary. That is, since merge can be invoked repeatedly without changing the results of the computation, a computation of the form:

@tff.federated_computation(...)
def single_aggregation(arg):
  result_at_clients = work(arg)
  agg_result = tff.federated_aggregate(
      result_at_clients, zero, accumulate, merge, report)
  return postprocess(arg, agg_result)

can be represented as the MergeableCompForm triplet:

@tff.federated_computation(tff.AbstractType('T'))
def up_to_merge(arg):
  result_at_clients = work(arg)
  agg_result = tff.federated_aggregate(
      result_at_clients, accumulate_zero, accumulate, merge, identity_report)
  return agg_result

@tff.federated_computation([up_to_merge.type_signature.result.member,
                            up_to_merge.type_signature.result.member])
def merge(arg):
  return merge(arg[0], arg[1])

@tff.federated_computation(
    tff.AbstractType('T'),
    tff.FederatedType(merge.type_signature.result, tff.SERVER),
)
def after_merge(original_arg, merged_result):
  return postprocess(original_arg, merged_result)

A fair amount of complexity can be hidden in postprocess; it could, for example, compute some value on clients based on the result of the aggregation. But the restriction that after_merge can contain no aggregations ensures that after_merge can also be executed in a subround fashion: a tff.CLIENTS-placed result can only depend on its own local state and the result of the aggregation, while a tff.SERVER-placed result can only depend on the result of the single aggregation or a tff.SERVER-placed value.