brainstate.environ module

`brainstate.environ` module#

Environment configuration and context management for BrainState.

This module provides comprehensive functionality for managing computational environments, including platform selection, precision control, mode setting, and context-based configuration management. It enables flexible configuration of JAX-based computations with thread-safe context switching.

The module supports: - Platform configuration (CPU, GPU, TPU) - Precision control (8, 16, 32, 64 bit and bfloat16) - Computation mode management - Context-based temporary settings - Default data type management - Custom behavior registration

Examples

Global environment configuration:

>>> import brainstate as bs
>>> import brainstate.environ as env
>>>
>>> # Set global precision to 32-bit
>>> env.set(precision=32, dt=0.01, mode=bs.mixin.Training())
>>>
>>> # Get current settings
>>> print(env.get('precision'))  # 32
>>> print(env.get('dt'))  # 0.01

Context-based temporary settings:

>>> import brainstate.environ as env
>>>
>>> # Temporarily change precision
>>> with env.context(precision=64, dt=0.001):
...     high_precision_result = compute_something()
...     print(env.get('precision'))  # 64
>>> print(env.get('precision'))  # Back to 32

The brainstate.environ module provides a comprehensive environment management system for configuring computational settings, platform preferences, numerical precision, and runtime behaviors. It offers thread-safe configuration management with support for both global settings and context-specific overrides.

Environment Management#

Functions for managing environment configuration and settings. These utilities allow you to set, retrieve, and manage environment variables that control the behavior of brainstate computations.

`reset`	Reset the environment to default settings.
`context`	Context manager for temporary environment settings.
`get`	Get a value from the current environment.
`set`	Set global environment configuration.
`pop`	Remove and return a value from the global environment.
`all`	Get all current environment settings.

Environment State#

Thread-local environment container used for custom or isolated configurations.

EnvironmentState

Thread-local storage for environment configuration.

Platform and Device Settings#

Configure the computing platform (CPU, GPU, TPU) and control device allocation. These functions help optimize performance by allowing you to select appropriate hardware backends and manage device resources.

`get_platform`	Get the current computing platform.
`set_platform`	Set the computing platform.
`get_host_device_count`	Get the number of host devices.
`set_host_device_count`	Set the number of host (CPU) devices.

Precision and Data Types#

Control numerical precision and retrieve appropriate data types for computations. These functions ensure consistent precision across your calculations and provide easy access to the correct NumPy/JAX data types based on the current precision setting.

`get_precision`	Get the current numerical precision as an integer.
`set_precision`	Set the global numerical precision.
`dftype`	Get the default floating-point data type.
`ditype`	Get the default integer data type.
`dutype`	Get the default unsigned integer data type.
`dctype`	Get the default complex data type.
`tolerance`	Get numerical tolerance based on current precision.

Mode and Timing#

Access computation mode settings and numerical integration parameters. These functions provide information about the current execution mode and time step settings for simulations.

get_dt

Get the current numerical integration time step.

Environment Keys and Flags#

Common environment key constants used with get/set/context.

`I`	str(object='') -> str str(bytes_or_buffer[, encoding[, errors]]) -> str
`T`	str(object='') -> str str(bytes_or_buffer[, encoding[, errors]]) -> str
`DT`	str(object='') -> str str(bytes_or_buffer[, encoding[, errors]]) -> str
`PRECISION`	str(object='') -> str str(bytes_or_buffer[, encoding[, errors]]) -> str
`PLATFORM`	str(object='') -> str str(bytes_or_buffer[, encoding[, errors]]) -> str
`HOST_DEVICE_COUNT`	str(object='') -> str str(bytes_or_buffer[, encoding[, errors]]) -> str
`JIT_ERROR_CHECK`	str(object='') -> str str(bytes_or_buffer[, encoding[, errors]]) -> str
`FIT`	str(object='') -> str str(bytes_or_buffer[, encoding[, errors]]) -> str

Defaults and Supported Values#

Default values and supported platform/precision settings.

`DEFAULT_PRECISION`	int([x]) -> integer int(x, base=10) -> integer
`SUPPORTED_PLATFORMS`	Built-in immutable sequence.
`SUPPORTED_PRECISIONS`	Built-in immutable sequence.

Behavior Registration#

Register custom callbacks that respond to environment parameter changes. This system allows you to define automatic behaviors when specific environment settings are modified, enabling reactive configuration management.

`register_default_behavior`	Register a callback for environment parameter changes.
`unregister_default_behavior`	Remove a registered callback for an environment parameter.
`list_registered_behaviors`	List all keys with registered callbacks.