@article{faxscalableanddifferentiablefederated,
  title = {DrJAX: Scalable and Differentiable MapReduce Primitives in JAX},
  author = {Keith Rush and Zachary Charles and Zachary Garrett and Sean Augenstein and Nicole Mitchell},
  year = {2024},
  eprint = {2403.07128},
  archivePrefix = {arXiv},
  url = {https://arxiv.org/abs/2403.07128v2},
}