I am experimenting with TensorFlow Federated and I would like to write a custom model aggregator for my training procedure, such that the distance of the client size from a certain predefined value is used to weight each client’s update (favouring some clients’ updates instead of others).

I would like to continue using the simple FedAvg algorithm that i have used in my previous simulations, to be able to compare the results:

trainer = tff.learning.algorithms.build_weighted_fed_avg(
                      model_fn= tff_model,
                      client_optimizer_fn=client_optimizer,
                      server_optimizer_fn=server_optimizer
                    )

I know that tff.learning.algorithms.build_weighted_fed_avg() accepts as a parameter a model_aggregator, but I don’t know how to create an aggregator that computes such weights.
Is there a simple way to define it?