py-mpi4jax (spack): Zero-copy MPI communication of JAX arrays, for turbo-charged HPC applications in Python. 
Keywords: gpu|high-performance-computing|jax|jit|mpi|parallel-computing|xla