flax

[pmap] Avoid degraded performance under the new `jax.pmap`.

#5152

Merged