flax
Make fully replicated array before saving checkpoints for examples that use pmap.
#4707
Merged
Merged