@article{fedjaxfederatedlearningsimulationwithjax,
  title = {FedJAX: Federated learning simulation with JAX},
  author = {Jae Hun Ro and Ananda Theertha Suresh and Ke wu},
  year = {2021},
  eprint = {2108.02117},
  archivePrefix = {arXiv},
  url = {https://arxiv.org/abs/2108.02117v2},
}