mpi4jax
MPI support for JAX
Installation
To install this package, run one of the following:
Usage Tracking
0.9.0
0.8.1.post2
0.8.1.post1
0.8.1
0.8.0
5 / 8 versions selected
Downloads (Last 6 months): 0
Description
Zero-copy MPI communication of JAX arrays, for turbo-charged HPC applications in Python ⚡
About
Summary
MPI support for JAX
Last Updated
Mar 5, 2026 at 14:30
License
MIT
Total Downloads
1.5M
Version Downloads
890
Supported Platforms
macOS-arm64
linux-64
Unsupported Platforms
macOS-64 Last supported version: 0.7.0
GitHub Repository
https://github.com/mpi4jax/mpi4jaxDocumentation
https://mpi4jax.readthedocs.io