set_key

Contents

set_key#

class brainstate.random.set_key(seed_or_key)#

Set a new random key for the global random state.

This function updates the global random state with a new key, which can be either an integer seed or a JAX PRNG key. All subsequent calls to global random functions will use this new key state.

Parameters:

seed_or_key (int | Array | ndarray) – The new random key to set. Can be: - An integer seed (will be converted to a JAX PRNG key) - A JAX PRNG key array - A numpy array representing a PRNG key

Raises:

ValueError – If the provided key is not in a valid format.

Return type:

None

Example

Set with integer seed:

>>> import brainstate
>>> brainstate.random.set_key(42)
>>> values1 = brainstate.random.rand(3)

Set with JAX key:

>>> import jax
>>> key = jax.random.key(123)
>>> brainstate.random.set_key(key)
>>> values2 = brainstate.random.rand(3)

Restore reproducible state:

>>> brainstate.random.set_key(42)
>>> # Now random functions will produce the same sequences as first example

Note

This function immediately changes the global random state. All threads and computations using the global random functions will be affected.

See also

Contents