def next_key(self) -> jax.Array:
    """Returns the next key in the sequence, and advances the stream."""
    old_offset = self.offset.value
    self.offset.value = self.offset.value + 1
    return jax.random.fold_in(self.base_key, old_offset)
Returns the next key in the sequence, and advances the stream.
penzai/core/random_stream.py