mlsanityData to 7 Oct 2026

DrJAX: Scalable and Differentiable MapReduce Primitives in JAX

arXiv 2403.07128 ↗First posted 11 Mar 2024cs.DC cs.LGcited by 9

Not on the map. Why?

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.

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

Builds on

26 of 39 references.

References · 26

  1. 2019PyTorch: An Imperative Style, High-Performance Deep Learning Librarycited by 20,160
  2. 2016Communication-Efficient Learning of Deep Networks from Decentralized Datacited by 7,592
  3. 2017Model-Agnostic Meta-Learning for Fast Adaptation of Deep Networkscited by 4,861
  4. 2020GShard: Scaling Giant Models with Conditional Computation and Automatic Shardingcited by 1,601
  5. 2022Model soups: averaging weights of multiple fine-tuned models improves accuracy without increasing inference timecited by 1,366
  6. 2015Automatic differentiation in machine learning: a surveycited by 846
  7. 2019Towards Federated Learning at Scale: System Designcited by 826
  8. 2020Flower: A Friendly Federated Learning Research Frameworkcited by 449
  9. 2022Branch-Train-Merge: Embarrassingly Parallel Training of Expert Language Modelscited by 187
  10. 2018Beyond Data and Model Parallelism for Deep Neural Networkscited by 177
  11. 2018Mesh-TensorFlow: Deep Learning for Supercomputerscited by 168
  12. 2021FedScale: Benchmarking Model and System Performance of Federated Learning at Scalecited by 135
14 more
  1. 2023DiLoCo: Distributed Low-Communication Training of Language Modelscited by 115
  2. 2021GSPMD: General and Scalable Parallelization for ML Computation Graphscited by 109
  3. 2021Papaya: Practical, Private, and Scalable Federated Learningcited by 89
  4. 2021Federated Evaluation and Tuning for On-Device Personalization: System Design & Applicationscited by 84
  5. 2022Pathways: Asynchronous Distributed Dataflow for MLcited by 81
  6. 2021FedJAX: Federated learning simulation with JAXcited by 32
  7. 2016Parallel SGD: When does averaging help?cited by 30
  8. 2023Towards Federated Foundation Models: Scalable Dataset Pipelines for Group-Structured Learningcited by 29
  9. 2022FLUTE: A Scalable, Extensible Framework for High-Performance Federated Learning Simulationscited by 27
  10. 2017Integrated Model, Batch and Domain Parallelism in Training Neural Networkscited by 16
  11. 2022Federated Select: A Primitive for Communication- and Memory-Efficient Federated Learningcited by 11
  12. 2022FL_PyTorch: optimization research simulator for federated learningcited by 10
  13. 2023FedHyper: A Universal and Robust Learning Rate Scheduler for Federated Learning with Hypergradient Descentcited by 8
  14. 2023Federated Automatic Differentiationcited by 2

Built on by

Cited by year (2026 to 7 Oct)

Citing papers per year: 2024: 1; 2025: 6; 2026: 2.

9 papers.

  1. 2024Fine-Tuning Large Language Models with User-Level Differential Privacycited by 28
  2. 2025Streaming DiLoCo with overlapping communication: Towards a Distributed Free Lunchcited by 26
  3. 2025Communication-Efficient Language Model Training Scales Reliably and Robustly: Scaling Laws for DiLoCocited by 23
  4. 2025Eager Updates For Overlapped Communication and Computation in DiLoCocited by 10
  5. 2025Understanding Outer Optimizers in Local SGD: Learning Rates, Momentum, and Accelerationcited by 7
  6. 2025Output Length Effect on DeepSeek-R1's Safety in Forced Thinkingcited by 6
  7. 2025Per-example gradients: a new frontier for understanding and improving optimizerscited by 2
  8. 2026A Typed Tensor Language for Shared-State Federated Computationcited by 0
  9. 2026Accelerator Choice Is Not Enough: AlphaFold2 Inference on Cloud TPUscited by 0