Rng

substrax.rng takes keys from an explicit owner and derives streams from a seed by name. Nothing in it falls back to a default seed: two callers never share a key they did not ask for.

from flax import nnx

from substrax.rng import fold_in_name, key_from, rngs_from_seed, split_key

rngs = rngs_from_seed(42, streams=("params", "dropout", "sample"))

# A sampler accepts an Rngs with a `sample` stream, or a key given directly.
key = key_from(rngs, streams=("sample", "default"), context="sampling")
per_walker = split_key(key, 1024)

# A sub-component derives its own key from a name, the same in every interpreter.
noise_key = fold_in_name(key, "observation_noise")

Function	Behaviour
key_from(rng, *, streams, context)	With an nnx.Rngs, calls the first of streams it holds, in the order given; a typed key or a raw uint32 key of shape (2,) is returned as given. A missing stream raises MissingRngStreamError naming context and the streams; None, a Python integer or any other array raises TypeError
rngs_from_seed(seed, streams=("default",))	One stream per name, each the seed's key folded by the name, so adding or reordering streams leaves the others unchanged
split_key(key, num)	jax.random.split with num at least one; asking for no keys is a defect
fold_in_name(key, name)	Folds in the name's BLAKE2b digest truncated to 32 bits, stable across interpreters whatever PYTHONHASHSEED is

rngs_from_seed differs from splitting the seed key into one part per stream, which every package used to do: there, appending a stream moved every existing stream's key, so a recorded run stopped reproducing when a model gained a dropout stream.

The pytest plugin (substrax.testing.pytest_plugin) provides rng_key, a typed key from seed 42, and rngs, an nnx.Rngs from the same seed with default, params, dropout and sample streams.

substrax.rng

Keys from an explicit owner, streams derived from a seed by name, splits and named folds.

key_from takes a key from an nnx.Rngs stream or a key given directly and never falls back to a default seed; rngs_from_seed derives each stream from the seed and the stream's name, so streams are independent of each other; split_key and fold_in_name are the two derivations every package needs, with a digest that is stable across interpreters.

MissingRngStreamError(context, streams)

Bases: ValueError

An nnx.Rngs holds none of the streams a call accepts.

Attributes:

Name	Type	Description
context	str	What needed the key, as the caller named it.
streams	tuple[str, ...]	The stream names the caller accepts, in order of preference.

Name what needed the key and the streams it accepts.

Parameters:

Name	Type	Description	Default
context	str	What needed the key, as the caller named it.	required
streams	Sequence[str]	The stream names the caller accepts, in order of preference.	required

fold_in_name(key, name)

Fold a name into key through a stable digest.

The name's BLAKE2b digest, truncated to 32 bits, is folded in, so the derived key is the same in every interpreter, whatever PYTHONHASHSEED is; Python's hash is salted per process and is never used.

Parameters:

Name	Type	Description	Default
key	Array	The key to derive from.	required
name	str	A non-empty name, such as a stream's.	required

Returns:

Type	Description
Array	The derived key.

Raises:

Type	Description
ValueError	If name is empty.

key_from(rng, *, streams, context)

Return a key from an explicit owner: an nnx.Rngs stream or a key given directly.

With an nnx.Rngs, the first of streams it holds is called, in the order given. A typed key (jax.random.key) or a raw uint32 key of shape (2,) is returned as given. There is no fallback: a missing stream, None, a Python integer or any other array raises, so two callers never share a key drawn from a default seed.

Parameters:

Name	Type	Description	Default
rng	Array | Rngs	The key's owner.	required
streams	Sequence[str]	Stream names accepted, most preferred first.	required
context	str	What needs the key, named in the errors ("sampling", "dropout").	required

Returns:

Type	Description
Array	The key.

Raises:

Type	Description
MissingRngStreamError	If rng is an nnx.Rngs holding none of streams.
TypeError	If rng is neither an nnx.Rngs nor a single key.

rngs_from_seed(seed, streams=('default',))

Build an nnx.Rngs whose every stream is the seed's key folded by the stream's name.

Each stream depends on the seed and its own name only, so adding or reordering streams leaves the others' keys unchanged, which a split of the seed key into len(streams) parts does not.

Parameters:

Name	Type	Description	Default
seed	int	A non-negative integer.	required
streams	Sequence[str]	Distinct stream names; at least one.	('default',)

Returns:

Type	Description
Rngs	The nnx.Rngs with one stream per name.

Raises:

Type	Description
TypeError	If seed is not an integer.
ValueError	If seed is negative, or streams is empty or repeats a name.

split_key(key, num)

Split key into num keys along a new leading axis.

Parameters:

Name	Type	Description	Default
key	Array	The key to split.	required
num	int	How many keys to return; at least one, so a caller asking for none is a defect.	required

Returns:

Type	Description
Array	Keys of shape (num,).

Raises:

Type	Description
ValueError	If num is below one.