flax

[pmap] In-line definitions of `jax.device_put_sharded` and `jax.device_put_replicated`.

#5374

Merged