def rng(self, rng):
    """Sets the rng (JAX PRNG key) for this layer and sublayers, recursively."""
    self._rng = rng
    sublayers = self.sublayers
    if sublayers:
      rngs = fastmath.random.split(rng, len(sublayers))
      for sublayer, rng in zip(sublayers, rngs):
        sublayer.rng = rng
Sets the rng (JAX PRNG key) for this layer and sublayers, recursively.
trax/layers/base.py