def next_comp(state, value, weight):
return collections.OrderedDict(
state=tff.federated_map(_add_one, state),
result=tff.federated_mean(value, weight),
measurements=tff.federated_zip(
collections.OrderedDict(num_clients=tff.federated_sum(
tff.federated_value(1, tff.CLIENTS)))))
def fed_output(local_outputs):
# TODO(b/124070381): Remove need for using num_examples_float here.
return {
'num_examples':
tff.federated_sum(local_outputs.num_examples),
'loss':
tff.federated_mean(local_outputs.loss,
weight=local_outputs.num_examples_float),
}
def federated_output_computation(self):
return tff.federated_computation(
lambda metrics: {'num_over': tff.federated_sum(metrics.num_over)})
def next_fn(empty_tup, x):
del empty_tup # Unused
return tff.federated_sum(tff.federated_map(reduce_dataset, x))
