custom aggregators with client_states as states - tensorflow-federated

Question

I want to create a custom aggregator where the state is the unique client state of each client. To initialize I can define client states as usual, and then use federated_collect to place #SERVER placement since thats what initialize_fn() wants. I can also do the same for creating new_state in the next_fn(). The problem is once I don't know how I can "broadcast" these states back into clients. Normally federated_broadcast takes say A#SERVER and then makes copies of it equal to number of clients. so for two clients it would be {A}#CLIENTS lets say (A A). What I want is to have AB#SERVER turning into (A B).
I am currently defining client states outside the aggregation process, and then passing these in run_one_round of iterative process. use federated_collect to collect these states from measurements of aggregator, and then unstack it outside. So from the outside of the federated computations it looks like
server_state, train_metrics , client_states, aggregation_state = iterative_process.next(
server_state, sampled_train_data, client_states, aggregation_state)
client_states = [x for x in client_states]
In TFF
output = aggregation_process.next(aggregation_state, client_outputs.weights_delta, client_states)
new_aggregation_state = output.state
round_model_delta = output.result
new_client_states = output.measurements
in aggregator
measurements = tff.federated_collect(new_client_states)
return tff.templates.MeasuredProcessOutput(
state=new_state, result= round_model_delta, measurements=measurements)
But I am trying to define and handle these client states completely inside the aggregator so that I can plug this aggregator into tff.learning.build_federated_averaging_process like
iterative_process = tff.learning.build_federated_averaging_process(
model_fn,
client_optimizer_fn=lambda: tf.keras.optimizers.SGD(learning_rate=0.02),
server_optimizer_fn=lambda: tf.keras.optimizers.SGD(learning_rate=1.0),
model_update_aggregation_factory=my_aggregation_factory)
Is that possible? if so how?

Answer 1

tff.federated_collect is likely not the desired tool in this situation, and it will be removed in future versions of TFF (see commit #030a406).
Alternatively, A tff.federated_computation can both take #CLIENTS parameters as input, and return #CLIENT placed values as output. Instead of collecting all the values on the server first (implying that the system is communicating the states); it maybe best to leave the values on the clients.
When executing TFF in a simulation environment (e.g. invoking a tff.Computation in a Colab notebook) a T#CLIENT placed value will be returned as a list of T objects; one for each client. This can later be used as a parameter to future tff.Computation invocation.
Example:
#tff.tf_computation(tf.int32)
def sqrt(value):
return tf.math.sqrt(tf.cast(value, tf.float32))
#tff.federated_computation(tff.types.at_clients(tf.int32))
def federated_sqrt(values):
return tff.federated_map(sqrt, values)
client_values = [1,2,3,4]
federated_sqrt(client_values)
>>> [<tf.Tensor: shape=(), dtype=float32, numpy=1.0>,
<tf.Tensor: shape=(), dtype=float32, numpy=1.4142135>,
<tf.Tensor: shape=(), dtype=float32, numpy=1.7320508>,
<tf.Tensor: shape=(), dtype=float32, numpy=2.0>]
Important caveat: the order of inputs and outputs is not necessarily guaranteed to be the same. An example of how to index and track states across invocations can be found in the tensorflow_federated/python/examples/stateful_clients/ directory inside the repository.