mpi4jax
mpi4jax enables zero-copy, multi-host communication of JAX arrays, even from jitted code and from GPU memory.
But why?
The JAX framework has great performance for scientific computing workloads, but its multi-host capabilities are still limited.
With mpi4jax, you can scale your JAX-based simulations to entire CPU and GPU clusters (without ever leaving jax.jit).
In the spirit of differentiable programming, mpi4jax also supports differentiating through some MPI operations.
Installation
mpi4jax is available through pip and conda:
$ pip install mpi4jax # Pip
$ conda install -c conda-forge mpi4jax # condaDepending on the different jax backends you want to use, you can install mpi4jax in the following way
# pip install 'jax[cpu]'
$ pip install mpi4jax
# pip install -U 'jax[cuda12]'
$ pip install nanobind
$ pip install mpi4jax --no-build-isolation
# pip install -U 'jax[cuda12_local]'
$ CUDA_ROOT=XXX pip install mpi4jax(for more informations on jax GPU distributions, see the JAX installation instructions)
In case your MPI installation is not detected correctly, it can help to install mpi4py separately. When using a pre-installed mpi4py, you must use --no-build-isolation when installing mpi4jax:
# if mpi4py is already installed
$ pip install cython
$ pip install mpi4jax --no-build-isolationOur documentation includes some more advanced installation examples.
Example usage
from mpi4py import MPI
import jax
import jax.numpy as jnp
import mpi4jax
comm = MPI.COMM_WORLD
rank = comm.Get_rank()
@jax.jit
def foo(arr):
arr = arr + rank
arr_sum, _ = mpi4jax.allreduce(arr, op=MPI.SUM, comm=comm)
return arr_sum
a = jnp.zeros((3, 3))
result = foo(a)
if rank == 0:
print(result)Running this script on 4 processes gives:
$ mpirun -n 4 python example.py
[[6. 6. 6.]
[6. 6. 6.]
[6. 6. 6.]]allreduce is just one example of the MPI primitives you can use. See all supported operations here.
Community guidelines
If you have a question or feature request, or want to report a bug, feel free to open an issue.
We welcome contributions of any kind through pull requests. For information on running our tests, debugging, and contribution guidelines please refer to the corresponding documentation page.
How to cite
If you use mpi4jax in your work, please consider citing the following article:
@article{mpi4jax,
doi = {10.21105/joss.03419},
url = {https://doi.org/10.21105/joss.03419},
year = {2021},
publisher = {The Open Journal},
volume = {6},
number = {65},
pages = {3419},
author = {Dion Häfner and Filippo Vicentini},
title = {mpi4jax: Zero-copy MPI communication of JAX arrays},
journal = {Journal of Open Source Software}
}
2 maintainers
Release files for mpi4jax 0.9.1.post2
For a detailed explanation of source distributions (sdists) and built distributions (wheels), please see the package formats documentation.
Source distribution (sdist)
| File | Size | Uploaded | |
|---|---|---|---|
| mpi4jax-0.9.1.post2.tar.gz | 65.5 kB | Details |
Release files / mpi4jax-0.9.1.post2.tar.gz
| Download URL | mpi4jax-0.9.1.post2.tar.gz |
|---|---|
| Size | 65.5 kB |
| Tags | Source |
|
SHA-256 checksum How to use checksums |
1089ab81b22edf7273e5c94cfac4dadb4fc4b2d691af87d6eff85479c8f15a0a
|
|
BLAKE2b-256 checksum How to use checksums |
2cecbc3513817a7216f2c2c14e36aa3fb77518c92c09e00d5215c6ef7d9c6ef8
|
| Upload date | |
|
Uploaded using Trusted Publishing? What is trusted publishing? |
Yes |
| Uploaded via |
twine/7.0.0 CPython/3.13.14
|
Provenance
Provenance describes where a file came from. On PyPI, provenance is shared via attestations, which provide a verifiable record of the build or publishing details. View details, limitations and caveats.
PyPI Publish Attestation
PyPI verified that this artifact, at this checksum, originated from the publisher listed below.
Signed by GitHub Actions, verified by PyPI on Sep 21, 2026.
Transparency log
