Revert change that breaks multi-host serialization with non-replicated data
Revert change that breaks multi-host serialization with non-replicated data

PiperOrigin-RevId: 470996705

google/flax,google/flax