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 server_init_tff():
"""Orchestration logic for server model initialization."""
@tff.tf_computation
def _fn():
return server_init(model_fn, server_optimizer_fn,
stateful_delta_aggregate_fn.initialize(),
stateful_model_broadcast_fn.initialize())
return tff.federated_value(_fn(), tff.SERVER)
def server_init_tff():
"""Orchestration logic for server model initialization."""
return tff.federated_value(tf_init_fn(), tff.SERVER)
def initialize_comp():
return tff.federated_value(0, tff.SERVER)
def server_init_tff(): """Orchestration logic for server model initialization.""" no_arg_server_init_fn = lambda: server_init(model_fn, server_optimizer_fn) server_init_tf = tff.tf_computation(no_arg_server_init_fn) return tff.federated_value(server_init_tf(), tff.SERVER)
def stateless_broadcast(state, value):
empty_metrics = tff.federated_value((), tff.SERVER)
return collections.OrderedDict(
state=state,
result=tff.federated_broadcast(value),
measurements=empty_metrics)
def stateless_mean(state, value, weight):
empty_metrics = tff.federated_value((), tff.SERVER)
return collections.OrderedDict(
state=state,
result=tff.federated_mean(value, weight=weight),
measurements=empty_metrics)
def _empty_server_initialization(): return tff.federated_value((), tff.SERVER)
def next_comp(state, value, weight):
empty_metrics = tff.federated_value((), tff.SERVER)
state, result = stateful_fn(state, value, weight)
return collections.OrderedDict(
state=state, result=result, measurements=empty_metrics)
