geossm.utils.KeyStream

class geossm.utils.KeyStream(seed)

Bases: object

Stateful helper that hands out fresh JAX PRNG keys on demand.

Wraps a single jax.random.PRNGKey and splits it on every call to next, so callers never have to thread/split keys manually or risk reusing the same key twice (e.g. across successive simulation calls in geossm.ssm.StateSpaceModel.sim()).

Parameters:

seed (jax.random.PRNGKey or int) – Initial key (or integer seed, though JAX generally expects an already-constructed PRNGKey) used to derive all subsequent keys.

Methods

next	Return one or more fresh keys, advancing the internal state.

next(num=None)

Return one or more fresh keys, advancing the internal state.

Parameters:

num (int, optional) – If given, return num independent keys (shape (num, ...)) instead of a single one; the internal state is advanced accordingly so the next call continues from a key independent of all of them.

Returns:

jax.random.PRNGKey or jax.numpy.ndarray – A single new key, or an array of num keys if num was given.