FedJAX: Federated learning simulation with JAX
  Federated learning is a machine learning technique that enables training
across decentralized data. Recently, federated learning has become an active
area of research due to an increased focus on privacy and security. In light of
this, a variety of open source federated learning libraries have been developed
and released. We introduce FedJAX, a JAX-based open source library for
federated learning simulations that emphasizes ease-of-use in research. With
its simple primitives for implementing federated learning algorithms,
prepackaged datasets, models and algorithms, and fast simulation speed, FedJAX
aims to make developing and evaluating federated algorithms faster and easier
for researchers. Our benchmark results show that FedJAX can be used to train
models with federated averaging on the EMNIST dataset in a few minutes and the
Stack Overflow dataset in roughly an hour with standard hyperparameters using
TPUs.
