flax

aa897790 - Make fully replicated array before saving checkpoints for examples that use pmap.

Make fully replicated array before saving checkpoints for examples that use pmap. This is required for multihost. Also removed pmap functions for examples that no longer use pmap (lm1b). PiperOrigin-RevId: 748461015