examples/research_projects/jax-projects/hybrid_clip/run_hybrid_clip.py/TrainState/replicate
class TrainState:    def replicate(self):
        return jax_utils.replicate(self).replace(dropout_rng=shard_prng_key(self.dropout_rng))