def new_rng(self):
    """Returns a new single-use random number generator (JAX PRNG key)."""
    self._rng, rng = fastmath.random.split(self._rng)
    if self._use_memory_efficient_trainer:
      self._rng = tl.on_cpu(self._rng)
      rng = tl.on_cpu(rng)
    return rng
Returns a new single-use random number generator (JAX PRNG key).
trax/supervised/training.py