seed_context

Contents

seed_context#

class brainstate.random.seed_context(seed_or_key)#

Context manager for temporary random seed changes with automatic restoration.

This context manager temporarily changes the global random seed for the duration of the block, then automatically restores the previous random state when exiting. It’s ideal for ensuring reproducible computations in specific code sections without permanently affecting the global random state.

Parameters:

seed_or_key (int | Array | ndarray) – The temporary seed or key to use within the context. Can be: - int: An integer seed for reproducible sequences - JAX PRNG key: A JAX random key array The seed affects both JAX and NumPy random states during the context.

Yields:

None – The context manager doesn’t yield any value, but provides a controlled random environment for the enclosed code block.

Examples

Reproducible computations without affecting global state:

>>> import brainstate
>>> # Global state remains unaffected
>>> global_values1 = brainstate.random.rand(2)
>>>
>>> with brainstate.random.seed_context(42):
...     temp_values1 = brainstate.random.rand(2)
...     print(f"First run: {temp_values1}")
[0.95598125 0.4032725 ]
>>>
>>> with brainstate.random.seed_context(42):
...     temp_values2 = brainstate.random.rand(2)
...     print(f"Second run: {temp_values2}")
[0.95598125 0.4032725 ]
>>>
>>> # Values are identical within context
>>> assert np.allclose(temp_values1, temp_values2)
>>>
>>> # Global state continues from where it left off
>>> global_values2 = brainstate.random.rand(2)

Nested contexts for complex scenarios:

>>> with brainstate.random.seed_context(123):
...     outer_values = brainstate.random.rand(2)
...     with brainstate.random.seed_context(456):
...         inner_values = brainstate.random.rand(2)
...     # Outer context is restored here
...     outer_values2 = brainstate.random.rand(2)

Exception safety - state is restored even on errors:

>>> try:
...     with brainstate.random.seed_context(789):
...         some_values = brainstate.random.rand(3)
...         raise ValueError("Something went wrong")
... except ValueError:
...     pass
>>> # Random state is properly restored

Testing reproducible algorithms:

>>> def test_algorithm():
...     with brainstate.random.seed_context(42):
...         data = brainstate.random.normal(size=(100,))
...         return data.mean()
>>>
>>> result1 = test_algorithm()
>>> result2 = test_algorithm()
>>> assert result1 == result2  # Always same result

Notes

  • The context manager saves and restores the complete JAX random state

  • NumPy’s global random state is also seeded on entry and fully restored on exit (via numpy.random.get_state/set_state), so NumPy-backed randomness inside the block is reproducible and leaves no side effect

  • Nested contexts work correctly - each level restores its own state

  • Exception safety is guaranteed - random state is restored even if exceptions occur within the context

  • This is more convenient than manually saving/restoring state with get_key() and set_key()

See also

seed

Permanently set the global random seed

get_key

Get the current random key for manual state management

set_key

Set the random key for manual state management

clone_rng

Create independent random states

Contents