flax
[pmap] Remove `jax.config.pmap_shmap_merge`.
#5289
Merged

Loading