LICENSE.md
MANIFEST.in
README.rst
pyproject.toml
setup.cfg
setup.py
versioneer.py
examples/shallow_water.py
mpi4jax/__init__.py
mpi4jax/_version.py
mpi4jax.egg-info/PKG-INFO
mpi4jax.egg-info/SOURCES.txt
mpi4jax.egg-info/dependency_links.txt
mpi4jax.egg-info/requires.txt
mpi4jax.egg-info/top_level.txt
mpi4jax/_src/__init__.py
mpi4jax/_src/_latest_jax_version.txt
mpi4jax/_src/comm.py
mpi4jax/_src/decorators.py
mpi4jax/_src/flush.py
mpi4jax/_src/jax_compat.py
mpi4jax/_src/utils.py
mpi4jax/_src/validation.py
mpi4jax/_src/collective_ops/__init__.py
mpi4jax/_src/collective_ops/allgather.py
mpi4jax/_src/collective_ops/allreduce.py
mpi4jax/_src/collective_ops/alltoall.py
mpi4jax/_src/collective_ops/barrier.py
mpi4jax/_src/collective_ops/bcast.py
mpi4jax/_src/collective_ops/gather.py
mpi4jax/_src/collective_ops/recv.py
mpi4jax/_src/collective_ops/reduce.py
mpi4jax/_src/collective_ops/scan.py
mpi4jax/_src/collective_ops/scatter.py
mpi4jax/_src/collective_ops/send.py
mpi4jax/_src/collective_ops/sendrecv.py
mpi4jax/_src/xla_bridge/__init__.py
mpi4jax/_src/xla_bridge/cuda_runtime_api.pxd
mpi4jax/_src/xla_bridge/mpi_xla_bridge.pxd
mpi4jax/_src/xla_bridge/mpi_xla_bridge.pyx
mpi4jax/_src/xla_bridge/mpi_xla_bridge_cpu.pyx
mpi4jax/_src/xla_bridge/mpi_xla_bridge_gpu.pyx
mpi4jax/experimental/__init__.py
mpi4jax/experimental/register_overrides.py
mpi4jax/experimental/tokenizer.py
tests/conftest.py
tests/test_decorators.py
tests/test_examples.py
tests/test_flush.py
tests/test_has_cuda.py
tests/test_jax_compat.py
tests/test_validation.py
tests/collective_ops/test_allgather.py
tests/collective_ops/test_allreduce.py
tests/collective_ops/test_allreduce_matvec.py
tests/collective_ops/test_alltoall.py
tests/collective_ops/test_barrier.py
tests/collective_ops/test_bcast.py
tests/collective_ops/test_common.py
tests/collective_ops/test_gather.py
tests/collective_ops/test_reduce.py
tests/collective_ops/test_scan.py
tests/collective_ops/test_scatter.py
tests/collective_ops/test_send_and_recv.py
tests/collective_ops/test_sendrecv.py
tests/experimental/test_auto_tokenize.py