I am implementing a client selection mechanism in TensorFlow Federated (TFF), where I need to dynamically choose the top_k clients with the largest dataset sizes before performing model update aggregation (deltas).

In TFF, tff.federated_map() executes functions independently on each client, meaning that each client only returns its own dataset size (sizes), and these values remain at the client level.