---
id: drjax
version: "0.2.0"
license: Apache License Version 2.0, January 2004 http://www.apache.org/licenses/ TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION 1. Definitions. "License" shall mean the terms and conditions for… (full text in the JSON record)
license_treatment: permissive
maintenance: active
---

# drjax — DrJAX - Scalable and Differentiable MapReduce Primitives in JAX.

License: permissive · Maintenance: active · Downloads: 253.4K/mo

## What it is and what it does

DrJAX is a JAX library that brings MapReduce-style distributed computing into the JAX ecosystem. It provides a simple authoring surface for parallel computations while leveraging JAX's sharding mechanisms for optimized execution in large-scale settings. The package is designed around three core goals: offering an accessible way to write MapReduce computations in JAX, enabling efficient execution via JAX's built-in parallelization, and maintaining full differentiability throughout—including through communication primitives like broadcasts and reductions.

The library embeds patterns similar to those in TensorFlow Federated but uses JAX's mapping and extension capabilities. It is tailored for scenarios involving large models and distributed training where both correctness and efficiency matter. The package depends on absl-py, jax, and jaxlib, and requires Python 3.9 or later. It is actively maintained and carries no known security vulnerabilities.

Use it for:

- Implement distributed training loops for large models using MapReduce patterns with automatic differentiation.
- Build federated learning workflows where computations are distributed across multiple devices or nodes.
- Optimize parallel data processing pipelines in datacenters by leveraging JAX's sharding for efficient communication.
- Prototype distributed algorithms that require differentiable reductions and broadcasts across compute clusters.
- Combine MapReduce-style data aggregation with JAX's automatic differentiation for gradient-based distributed optimization.

## Worth the install?

AI-flagged interpretation of the facts above — verify before relying.

DrJAX embeds MapReduce programming primitives into JAX, enabling differentiable distributed and parallel computations with optimized execution across large-scale datacenter settings.

Yes, if you are building distributed or federated machine learning systems in JAX and need MapReduce-style primitives with automatic differentiation. The package is actively maintained, carries no known vulnerabilities, and has low install friction. However, it remains early-stage (first release January 2025) with modest adoption; verify that its API and performance characteristics match your production requirements before committing to it.

## Install

pip install drjax
uv add drjax
poetry add drjax

## Installing drjax

Before you install: Low friction installation with a pure Python wheel. The package is actively maintained with a recent release and no known vulnerabilities, though it remains early-stage with modest adoption.

License in practice: Apache 2.0 permissive license allows commercial and derivative use with minimal restrictions; you must include license and copyright notices in distributions.

Quickstart:

pip install drjax

import drjax
import jax

# Define a MapReduce computation using DrJAX primitives
# and execute it with JAX's sharding mechanisms

Requires Python 3.9 or later; jax and jaxlib must be installed as runtime dependencies.

Verify before relying:

- Whether the package's MapReduce API is stable or subject to breaking changes in future releases.
- Performance characteristics and scalability limits in typical datacenter deployments.
- Availability of tutorials or examples beyond the research paper for practical adoption.

## Package facts

- License: Apache License Version 2.0, January 2004 http://www.apache.org/licenses/ TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION 1. Definitions. "License" shall mean the terms and conditions for… (full text in the JSON record) (permissive)
- Python support: supports_current
- Install friction: low
- Maintenance: active
- Downloads: 253.4K/month (top 15,000 on PyPI)
- Known vulnerabilities: none known

## Tags

mapreduce jax, distributed parallel computing jax, differentiable mapreduce, jax distributed primitives, datacenter scale computation, jax sharding optimization, federated learning jax, distributed-computing, jax-ecosystem, federated-learning

[View on SkillFed](https://skillfed.io/packages/drjax) · [View on PyPI](https://pypi.org/project/drjax/)