DrJAX: Scalable and Differentiable MapReduce Primitives in JAX

Understand

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.