def from_base_key(
      cls,
      base_key: jax.Array,
      offset_label: variables.VariableLabel = "random_stream_offset",
  ) -> RandomStream:
    """Returns a new random stream with the given base key."""
    return cls(
        base_key=base_key, offset=variables.StateVariable(0, label=offset_label)
    )
Returns a new random stream with the given base key.
penzai/core/random_stream.py