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 an integer seed, a JAX typed PRNG key, or a legacy uint32[2] key array (auto-wrapped into a typed 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 (converted to a typed JAX key via jax.random.key) - A JAX typed PRNG key - A legacy uint32[2] array (auto-wrapped into a typed key)

Raises:

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

Return type:

None

Examples

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

Notes

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

See also

get_key

Get the current random key

get_key_data

Get the current key as raw uint32[2] data

seed

Set seed (also affects NumPy)

restore_key

Restore a backed up key

Contents