Parallax

Parallax is a library for easily scaling off-the-shelf JAX models. It includes optimizations in two main categories:

• Parallelism:
  • AutoSharding: Automatically determine PartitionSpecs for JAX Models and Functions.
  • Sharding Recipes: Apply Distributed Data Parallel (DDP) or Fully Sharded Data Parallel (FSDP) to JAX models.
• Memory optimizations:
  • Model Offloading: Keep model weights in CPU RAM and stream to TPU/GPU HBM as needed.
  • Rematerialization: Intentionally recompute activations to save HBM space.

The ethos of Parallax is humans shouldn't be writing their own Mesh or PartitionSpec — computers are much better at discovering optimal sharding strategies, so let them do the hard work.

This is not an officially supported Google product. This project is not eligible for the Google Open Source Software Vulnerability Rewards Program.

Why Parallax?

With plain JAX, you write sharding rules for every parameter by hand. Parallax does that work for you and lets you take over where you need to:

Task	Manual JAX	Parallax
Choose parameter shardings	Write PartitionSpec rules by hand	strategy=parallax.ShardingStrategy.AUTO
Apply DDP or FSDP	Write sharding rules for each model	strategy=parallax.ShardingStrategy.DDP (or .FSDP)
Rematerialize activations	Wrap layers with jax.checkpoint	parallax.remat_model(model) (sequential models)
Offload weights	Manage memory kinds and transfers	parallax.offload_model(model, ...)

Parallax is most useful for large models, or when you want to try different parallelism strategies without rewriting your model.

Install Parallax

Install the latest release of Parallax from PyPI:

pip install google-parallax

Usage

Parallax integrates seamlessly with Flax NNX models. Here is a simple example of how to use auto-sharding:

import jax
import jax.numpy as jnp
from flax import nnx

import parallax

with jax.set_mesh(mesh):  # a 2D mesh with 'fsdp' and 'tp' axes
  model = parallax.create_sharded_model(
      model_or_fn=lambda: Model(...),
      sample_inputs=(jnp.ones((1, 32)),),
      strategy=parallax.ShardingStrategy.AUTO,
      data_axis_name='fsdp',
      model_axis_name='tp',
  )

For a complete, runnable example, see the Getting Started notebook.

Support

For questions, bug reports, or feature requests, please open an issue on GitHub Issues.

Contribute

To contribute to Parallax, see the contributing guide.

Citing Parallax

To cite Parallax please use this citation:

@software{parallax2026,
  title = {{Parallax}},
  author = {Carpenter, Jeff and Yan, Stella and Liu, Dangyi and others},
  year = {2026},
  url = {https://github.com/google/parallax},
}