FR: Create mpi4jax.pmap working as a GlobalShardedDeviceArray #122
Description
We could implement a tracer, mpi4jax.pmap, which behaves as jax.pmap but is multi-process aware.
{{ message }}
We could implement a tracer, mpi4jax.pmap, which behaves as jax.pmap but is multi-process aware.