flax

[graphcast] Prepare for `jax_pmap_shmap_merge=True`.

#5168

Open