Make sure save checkpoint only happens on process_index 0. This will make save_checkpoint truly multihost-friendly. Also made some refactoring & extra commenting.
Make sure save checkpoint only happens on process_index 0. This will make save_checkpoint truly multihost-friendly. Also made some refactoring & extra commenting.

PiperOrigin-RevId: 470852413

google/flax,google/flax