def next_key(self) -> jax.Array:
    """Gets the next key from this stream, mutating the stream in place."""
    if self.state != "active":
      if self.state == "pending":
        raise ValueError(
            "Cannot use a random stream that is not yet active! For safety,"
            " random streams should be activated by using them as a context"
            " manager, e.g. `with RandomStream(key) as stream:`, and they are"
            " only active within that context.\n\nIf you want to use a random"
            " stream at the top level (e.g. in a notebook), you can call"
            " `.unsafe_mark_active()`. However, in that case, make sure not to"
            " use this random stream inside a function being transformed by"
            " JAX, since this may lead to unexpected results."
        )
      elif self.state == "expired":
        raise ValueError(
            "Cannot use a random stream that has expired! For safety, random"
            " streams are only active within a scope defined by a context"
            " manager, e.g. `with RandomStream(key) as stream:`. This random"
            " stream has now been accessed outside the context it was"
            " defined in, which is not allowed."
        )
      else:
        raise ValueError(f"Unexpected state for RandomStream: {self.state}")
    offset = self.next_offset
    self.next_offset += 1
    return jax.random.fold_in(self.base_key, offset)
Gets the next key from this stream, mutating the stream in place.
penzai/deprecated/v1/core/random_stream.py