set

Contents

set#

class brainstate.environ.set(platform=None, host_device_count=None, precision=None, dt=None, *, env=None, **kwargs)[source]#

Set global environment configuration.

This function sets persistent global environment settings that remain active until explicitly changed or the program terminates.

Parameters:

  • platform (str | None) – Computing platform (‘cpu’, ‘gpu’, or ‘tpu’).

  • host_device_count (int | None) – Number of host devices for parallel computation.

  • precision (int | str | None) – Numerical precision (8, 16, 32, 64, or ‘bf16’).

  • mode (Mode, optional) – Computation mode instance.

  • dt (float | None) – Time step for numerical integration.

  • env (EnvironmentState | None) – The environment state to modify. If None, uses the global environment.

  • **kwargs – Additional custom environment parameters.

Raises:

  • ValueError – If invalid platform or precision is specified.

  • TypeError – If mode is not a Mode instance.

Return type:

None

Examples

Basic configuration:

>>> import brainstate as bs
>>> import brainstate.environ as env
>>>
>>> # Set multiple parameters
>>> env.set(
...     precision=32,
...     dt=0.01,
...     mode=bs.mixin.Training(),
...     debug=False
... )
>>>
>>> print(env.get('precision'))  # 32
>>> print(env.get('dt'))  # 0.01

Platform configuration:

>>> import brainstate.environ as env
>>>
>>> # Configure for GPU computation
>>> env.set(platform='gpu', precision=16)
>>>
>>> # Configure for multi-core CPU
>>> env.set(platform='cpu', host_device_count=4)

Custom parameters:

>>> import brainstate.environ as env
>>>
>>> # Set custom parameters
>>> env.set(
...     experiment_name='test_001',
...     random_seed=42,
...     log_level='DEBUG'
... )
>>>
>>> # Retrieve custom parameters
>>> print(env.get('experiment_name'))  # 'test_001'

Using custom environment:

>>> import brainstate.environ as env
>>>
>>> custom_env = env.EnvironmentState()
>>> env.set(precision=64, dt=0.001, env=custom_env)
>>> print(env.get('precision', env=custom_env))  # 64

Notes

  • Platform changes only take effect at program start

  • Some JAX configurations require restart to take effect

  • Custom parameters can be any hashable key-value pairs

  • When using a custom env, JAX config is only updated if env is the global environment

Contents