drjax

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 on this page — 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

drjax on PyPI

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

Evidence: drjax-0.2.0-py3-none-any.whl

Keywords: parallel, mapreduce, machine, learning, jax

Tags

More Distributed Computing packages