def _per_device_rngs(self, rng):
    """Create per-device RNGs from a given rng."""
    # Splitting by device first to be identical with default trainer.
    per_device_rng = fastmath.random.split(rng, self._n_devices)
    per_device_rngs = [
        fastmath.random.split(r, self._n_layers) for r in per_device_rng]
    rngs = [jnp.stack([r[i] for r in per_device_rngs])
            for i in range(self._n_layers)]
    return rngs
Create per-device RNGs from a given rng.
trax/optimizers/trainer.py