# DrJAX: Scalable and Differentiable MapReduce Primitives in JAX

arXiv: https://arxiv.org/abs/2403.07128 · HTML: https://mlsanity.com/paper/2403.07128/

## Key facts

- Cited by 9 papers in the map (data through 14 Oct 2026).
- Builds on 26 of its 39 references that are in the map.
- First posted on arXiv 11 Mar 2024.

## Builds on (most cited first, 26 of 26)

- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [Communication-Efficient Learning of Deep Networks from Decentralized Data](https://mlsanity.com/paper/1602.05629/) (2016)
- [Model-Agnostic Meta-Learning for Fast Adaptation of Deep Networks](https://mlsanity.com/paper/1703.03400/) (2017)
- [GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding](https://mlsanity.com/paper/2006.16668/) (2020)
- [Model soups: averaging weights of multiple fine-tuned models improves accuracy without increasing inference time](https://mlsanity.com/paper/2203.05482/) (2022)
- [Automatic differentiation in machine learning: a survey](https://mlsanity.com/paper/1502.05767/) (2015)
- [Towards Federated Learning at Scale: System Design](https://mlsanity.com/paper/1902.01046/) (2019)
- [Flower: A Friendly Federated Learning Research Framework](https://mlsanity.com/paper/2007.14390/) (2020)
- [Branch-Train-Merge: Embarrassingly Parallel Training of Expert Language Models](https://mlsanity.com/paper/2208.03306/) (2022)
- [Beyond Data and Model Parallelism for Deep Neural Networks](https://mlsanity.com/paper/1807.05358/) (2018)
- [Mesh-TensorFlow: Deep Learning for Supercomputers](https://mlsanity.com/paper/1811.02084/) (2018)
- [FedScale: Benchmarking Model and System Performance of Federated Learning at Scale](https://mlsanity.com/paper/2105.11367/) (2021)
- [DiLoCo: Distributed Low-Communication Training of Language Models](https://mlsanity.com/paper/2311.08105/) (2023)
- [GSPMD: General and Scalable Parallelization for ML Computation Graphs](https://mlsanity.com/paper/2105.04663/) (2021)
- [Papaya: Practical, Private, and Scalable Federated Learning](https://mlsanity.com/paper/2111.04877/) (2021)
- [Federated Evaluation and Tuning for On-Device Personalization: System Design & Applications](https://mlsanity.com/paper/2102.08503/) (2021)
- [Pathways: Asynchronous Distributed Dataflow for ML](https://mlsanity.com/paper/2203.12533/) (2022)
- [Parallel SGD: When does averaging help?](https://mlsanity.com/paper/1606.07365/) (2016)
- [FedJAX: Federated learning simulation with JAX](https://mlsanity.com/paper/2108.02117/) (2021)
- [Towards Federated Foundation Models: Scalable Dataset Pipelines for Group-Structured Learning](https://mlsanity.com/paper/2307.09619/) (2023)
- [FLUTE: A Scalable, Extensible Framework for High-Performance Federated Learning Simulations](https://mlsanity.com/paper/2203.13789/) (2022)
- [Integrated Model, Batch and Domain Parallelism in Training Neural Networks](https://mlsanity.com/paper/1712.04432/) (2017)
- [Federated Select: A Primitive for Communication- and Memory-Efficient Federated Learning](https://mlsanity.com/paper/2208.09432/) (2022)
- [FL_PyTorch: optimization research simulator for federated learning](https://mlsanity.com/paper/2202.03099/) (2022)
- [FedHyper: A Universal and Robust Learning Rate Scheduler for Federated Learning with Hypergradient Descent](https://mlsanity.com/paper/2310.03156/) (2023)
- [Federated Automatic Differentiation](https://mlsanity.com/paper/2301.07806/) (2023)

## Built on by (most cited first, 9 of 9)

- [Fine-Tuning Large Language Models with User-Level Differential Privacy](https://mlsanity.com/paper/2407.07737/) (2024)
- [Streaming DiLoCo with overlapping communication: Towards a Distributed Free Lunch](https://mlsanity.com/paper/2501.18512/) (2025)
- [Communication-Efficient Language Model Training Scales Reliably and Robustly: Scaling Laws for DiLoCo](https://mlsanity.com/paper/2503.09799/) (2025)
- [Eager Updates For Overlapped Communication and Computation in DiLoCo](https://mlsanity.com/paper/2502.12996/) (2025)
- [Understanding Outer Optimizers in Local SGD: Learning Rates, Momentum, and Acceleration](https://mlsanity.com/paper/2509.10439/) (2025)
- [Output Length Effect on DeepSeek-R1's Safety in Forced Thinking](https://mlsanity.com/paper/2503.01923/) (2025)
- [Per-example gradients: a new frontier for understanding and improving optimizers](https://mlsanity.com/paper/2510.00236/) (2025)
- [A Typed Tensor Language for Shared-State Federated Computation](https://mlsanity.com/paper/2605.21103/) (2026)
- [Accelerator Choice Is Not Enough: AlphaFold2 Inference on Cloud TPUs](https://mlsanity.com/paper/2609.34818/) (2026)

## Abstract (arXiv)

We present DrJAX, a JAX-based library designed to support large-scale distributed and parallel machine learning algorithms that use MapReduce-style operations. DrJAX leverages JAX's sharding mechanisms to enable native targeting of TPUs and state-of-the-art JAX runtimes, including Pathways. DrJAX embeds building blocks for MapReduce computations as primitives in JAX. This enables three key benefits. First, DrJAX computations can be translated directly to XLA HLO, enabling flexible integration with a wide array of ML training platforms. Second, DrJAX computations are fully differentiable. Last, DrJAX computations can be interpreted out to existing batch-processing compute systems, including traditional MapReduce systems like Apache Beam and cross-device compute systems like those powering federated learning applications. We show that DrJAX provides an easily programmable, performant, and scalable framework for parallelized algorithm development. DrJAX is available at urlhttps://github.com/google-research/google-research/tree/master/drjax.