brainstate 0.1.9__py2.py3-none-any.whl → 0.2.0__py2.py3-none-any.whl

brainstate/environ.py

@@ -1,4 +1,4 @@
-# Copyright 2024 BDP Ecosystem Limited. All Rights Reserved.
+# Copyright 2024 BrainX Ecosystem Limited. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -13,7 +13,50 @@
 # limitations under the License.
 # ==============================================================================
 
-# -*- coding: utf-8 -*-
+"""
+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:
+
+.. code-block:: python
+
+    >>> 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:
+
+.. code-block:: python
+
+    >>> 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
+"""
 
 import contextlib
 import dataclasses
@@ -21,543 +64,1432 @@ import functools
 import os
 import re
 import threading
+import warnings
 from collections import defaultdict
-from typing import Any, Callable, Dict, Hashable
+from typing import Any, Callable, Dict, Hashable, Optional, Union, ContextManager, List
 
+import brainunit as u
 import numpy as np
 from jax import config, devices, numpy as jnp
 from jax.typing import DTypeLike
 
-from .mixin import Mode
-
 __all__ = [
-    # functions for environment settings
-    'set', 'context', 'get', 'all', 'set_host_device_count', 'set_platform',
-    # functions for getting default behaviors
-    'get_host_device_count', 'get_platform', 'get_dt', 'get_mode', 'get_precision',
-    # functions for default data types
-    'dftype', 'ditype', 'dutype', 'dctype',
-    # others
-    'tolerance', 'register_default_behavior',
+    # Core environment management
+    'set',
+    'get',
+    'all',
+    'pop',
+    'context',
+    'reset',
+
+    # Platform and device management
+    'set_platform',
+    'get_platform',
+    'set_host_device_count',
+    'get_host_device_count',
+
+    # Precision and data type management
+    'set_precision',
+    'get_precision',
+    'dftype',
+    'ditype',
+    'dutype',
+    'dctype',
+
+    # Mode and computation settings
+    'get_dt',
+
+    # Utility functions
+    'tolerance',
+    'register_default_behavior',
+    'unregister_default_behavior',
+    'list_registered_behaviors',
+
+    # Constants
+    'DEFAULT_PRECISION',
+    'SUPPORTED_PLATFORMS',
+    'SUPPORTED_PRECISIONS',
+
+    # Names
+    'I',
+    'T',
+    'DT',
+    'PRECISION',
+    'PLATFORM',
+    'HOST_DEVICE_COUNT',
+    'JIT_ERROR_CHECK',
+    'FIT',
 ]
 
-# Default, there are several shared arguments in the global context.
-I = 'i'  # the index of the current computation.
-T = 't'  # the current time of the current computation.
-JIT_ERROR_CHECK = 'jit_error_check'  # whether to record the current computation.
-FIT = 'fit'  # whether to fit the model.
+# Type definitions
+# T = TypeVar('T')
+PrecisionType = Union[int, str]
+PlatformType = str
 
+# Constants for environment keys
+I = 'i'  # Index of the current computation
+T = 't'  # Current time of the computation
+DT = 'dt'  # Time step for numerical integration
+PRECISION = 'precision'  # Numerical precision
+PLATFORM = 'platform'  # Computing platform
+HOST_DEVICE_COUNT = 'host_device_count'  # Number of host devices
+JIT_ERROR_CHECK = 'jit_error_check'  # JIT error checking flag
+FIT = 'fit'  # Model fitting flag
 
-@dataclasses.dataclass
-class DefaultContext(threading.local):
-    # default environment settings
-    settings: Dict[Hashable, Any] = dataclasses.field(default_factory=dict)
-    # current environment settings
-    contexts: defaultdict[Hashable, Any] = dataclasses.field(default_factory=lambda: defaultdict(list))
-    # environment functions
-    functions: Dict[Hashable, Any] = dataclasses.field(default_factory=dict)
+# Default values
+DEFAULT_PRECISION = 32
+SUPPORTED_PLATFORMS = ('cpu', 'gpu', 'tpu')
+SUPPORTED_PRECISIONS = (8, 16, 32, 64, 'bf16')
 
+# Sentinel value for missing arguments
+_NOT_PROVIDED = object()
 
-DFAULT = DefaultContext()
-_NOT_PROVIDE = object()
 
+@dataclasses.dataclass
+class EnvironmentState(threading.local):
+    """
+    Thread-local storage for environment configuration.
 
-@contextlib.contextmanager
-def context(**kwargs):
-    r"""
-    Context-manager that sets a computing environment for brain dynamics computation.
+    This class maintains separate configuration states for different threads,
+    ensuring thread-safe environment management in concurrent applications.
 
-    In BrainPy, there are several basic computation settings when constructing models,
-    including ``mode`` for controlling model computing behavior, ``dt`` for numerical
-    integration, ``int_`` for integer precision, and ``float_`` for floating precision.
-    :py:class:`~.environment`` provides a context for model construction and
-    computation. In this temporal environment, models are constructed with the given
-    ``mode``, ``dt``, ``int_``, etc., environment settings.
+    Attributes
+    ----------
+    settings : Dict[Hashable, Any]
+        Global default environment settings.
+    contexts : defaultdict[Hashable, List[Any]]
+        Stack of context-specific settings for nested contexts.
+    functions : Dict[Hashable, Callable]
+        Registered callback functions for environment changes.
+    locks : Dict[str, threading.Lock]
+        Thread locks for synchronized access to critical sections.
+    """
+    settings: Dict[Hashable, Any] = dataclasses.field(default_factory=dict)
+    contexts: defaultdict[Hashable, List[Any]] = dataclasses.field(default_factory=lambda: defaultdict(list))
+    functions: Dict[Hashable, Callable] = dataclasses.field(default_factory=dict)
+    locks: Dict[str, threading.Lock] = dataclasses.field(default_factory=lambda: defaultdict(threading.Lock))
 
-    For instance::
+    def __post_init__(self):
+        """Initialize with default settings."""
+        # Set default precision if not already set
+        if PRECISION not in self.settings:
+            self.settings[PRECISION] = DEFAULT_PRECISION
 
-    >>> import brainstate as brainstate
-    >>> with brainstate.environ.context(dt=0.1) as env:
-    ...     dt = brainstate.environ.get('dt')
-    ...     print(env)
 
-    """
-    if 'platform' in kwargs:
-        raise ValueError('\n'
-                         'Cannot set platform in "context" environment. \n'
-                         'You should set platform in the global environment by "set_platform()" or "set()".')
-    if 'host_device_count' in kwargs:
-        raise ValueError('Cannot set host_device_count in environment context. '
-                         'Please use set_host_device_count() or set() for the global setting.')
+# Global environment state
+_ENV_STATE = EnvironmentState()
 
-    if 'precision' in kwargs:
-        last_precision = _get_precision()
-        _set_jax_precision(kwargs['precision'])
 
-    try:
-        for k, v in kwargs.items():
+def reset() -> None:
+    """
+    Reset the environment to default settings.
+
+    This function clears all custom settings and restores the environment
+    to its initial state. Useful for testing or when starting fresh.
+
+    Examples
+    --------
+    .. code-block:: python
+
+        >>> import brainstate.environ as env
+        >>>
+        >>> # Set custom values
+        >>> env.set(dt=0.1, custom_param='value')
+        >>> print(env.get('custom_param'))  # 'value'
+        >>>
+        >>> # Reset to defaults
+        >>> env.reset()
+        >>> print(env.get('custom_param', default=None))  # None
+
+    Notes
+    -----
+    This operation cannot be undone. All custom settings will be lost.
+    """
+    global _ENV_STATE
+    _ENV_STATE = EnvironmentState()
+    # Re-apply default precision
+    _set_jax_precision(DEFAULT_PRECISION)
 
-            # update the current environment
-            DFAULT.contexts[k].append(v)
+    warnings.warn(
+        "Environment has been reset to default settings. "
+        "All custom configurations have been cleared.",
+        UserWarning
+    )
 
-            # restore the environment functions
-            if k in DFAULT.functions:
-                DFAULT.functions[k](v)
 
-        # yield the current all environment information
-        yield all()
-    finally:
+@contextlib.contextmanager
+def context(**kwargs) -> ContextManager[Dict[str, Any]]:
+    """
+    Context manager for temporary environment settings.
 
-        for k, v in kwargs.items():
+    This context manager allows you to temporarily modify environment settings
+    within a specific scope. Settings are automatically restored when exiting
+    the context, even if an exception occurs.
 
-            # restore the current environment
-            DFAULT.contexts[k].pop()
+    Parameters
+    ----------
+    **kwargs
+        Environment settings to apply within the context.
+        Common parameters include:
+
+        - precision : int or str.
+            Numerical precision (8, 16, 32, 64, or 'bf16')
+        - dt : float.
+            Time step for numerical integration
+        - mode : :class:`Mode`.
+            Computation mode instance
+        - Any custom parameters registered via register_default_behavior
+
+    Yields
+    ------
+    dict
+        Current environment settings within the context.
+
+    Raises
+    ------
+    ValueError
+        If attempting to set platform or host_device_count in context
+        (these must be set globally).
+    TypeError
+        If invalid parameter types are provided.
+
+    Examples
+    --------
+    Basic usage with precision control:
+
+    .. code-block:: python
+
+        >>> import brainstate.environ as env
+        >>>
+        >>> # Set global precision
+        >>> env.set(precision=32)
+        >>>
+        >>> # Temporarily use higher precision
+        >>> with env.context(precision=64) as ctx:
+        ...     print(f"Precision in context: {env.get('precision')}")  # 64
+        ...     print(f"Float type: {env.dftype()}")  # float64
+        >>>
+        >>> print(f"Precision after context: {env.get('precision')}")  # 32
+
+    Nested contexts:
+
+    .. code-block:: python
+
+        >>> import brainstate.environ as env
+        >>>
+        >>> with env.context(dt=0.1) as ctx1:
+        ...     print(f"dt = {env.get('dt')}")  # 0.1
+        ...
+        ...     with env.context(dt=0.01) as ctx2:
+        ...         print(f"dt = {env.get('dt')}")  # 0.01
+        ...
+        ...     print(f"dt = {env.get('dt')}")  # 0.1
+
+    Error handling in context:
+
+    .. code-block:: python
+
+        >>> import brainstate.environ as env
+        >>>
+        >>> env.set(value=10)
+        >>> try:
+        ...     with env.context(value=20):
+        ...         print(env.get('value'))  # 20
+        ...         raise ValueError("Something went wrong")
+        ... except ValueError:
+        ...     pass
+        >>>
+        >>> print(env.get('value'))  # 10 (restored)
+
+    Notes
+    -----
+    - Platform and host_device_count cannot be set in context
+    - Contexts can be nested arbitrarily deep
+    - Settings are restored in reverse order when exiting
+    - Thread-safe: each thread maintains its own context stack
+    """
+    # Validate restricted parameters
+    if PLATFORM in kwargs:
+        raise ValueError(
+            f"Cannot set '{PLATFORM}' in context. "
+            f"Use set_platform() or set() for global configuration."
+        )
+    if HOST_DEVICE_COUNT in kwargs:
+        raise ValueError(
+            f"Cannot set '{HOST_DEVICE_COUNT}' in context. "
+            f"Use set_host_device_count() or set() for global configuration."
+        )
+
+    # Handle precision changes
+    original_precision = None
+    if PRECISION in kwargs:
+        original_precision = _get_precision()
+        _validate_precision(kwargs[PRECISION])
+        _set_jax_precision(kwargs[PRECISION])
 
-            # restore the environment functions
-            if k in DFAULT.functions:
-                DFAULT.functions[k](get(k))
+    try:
+        # Push new values onto context stacks
+        for key, value in kwargs.items():
+            with _ENV_STATE.locks[key]:
+                _ENV_STATE.contexts[key].append(value)
+
+                # Trigger registered callbacks
+                if key in _ENV_STATE.functions:
+                    try:
+                        _ENV_STATE.functions[key](value)
+                    except Exception as e:
+                        warnings.warn(
+                            f"Callback for '{key}' raised an exception: {e}",
+                            RuntimeWarning
+                        )
+
+        # Yield current environment state
+        yield all()
 
-        if 'precision' in kwargs:
-            _set_jax_precision(last_precision)
+    finally:
+        # Restore previous values
+        for key in kwargs:
+            with _ENV_STATE.locks[key]:
+                if _ENV_STATE.contexts[key]:
+                    _ENV_STATE.contexts[key].pop()
+
+                # Restore callbacks with previous value
+                if key in _ENV_STATE.functions:
+                    try:
+                        prev_value = get(key, default=None)
+                        if prev_value is not None:
+                            _ENV_STATE.functions[key](prev_value)
+                    except Exception as e:
+                        warnings.warn(
+                            f"Callback restoration for '{key}' raised: {e}",
+                            RuntimeWarning
+                        )
+
+        # Restore precision if it was changed
+        if original_precision is not None:
+            _set_jax_precision(original_precision)
+
+
+def get(key: str, default: Any = _NOT_PROVIDED, desc: Optional[str] = None) -> Any:
+    """
+    Get a value from the current environment.
 
+    This function retrieves values from the environment, checking first in
+    the context stack, then in global settings. Special handling is provided
+    for platform and device count parameters.
 
-def get(key: str, default: Any = _NOT_PROVIDE, desc: str = None):
-    """
-    Get one of the default computation environment.
+    Parameters
+    ----------
+    key : str
+        The environment key to retrieve.
+    default : Any, optional
+        Default value to return if key is not found.
+        If not provided, raises KeyError for missing keys.
+    desc : str, optional
+        Description of the parameter for error messages.
 
     Returns
     -------
-    item: Any
-      The default computation environment.
+    Any
+        The value associated with the key.
+
+    Raises
+    ------
+    KeyError
+        If key is not found and no default is provided.
+
+    Examples
+    --------
+    Basic retrieval:
+
+    .. code-block:: python
+
+        >>> import brainstate.environ as env
+        >>>
+        >>> env.set(learning_rate=0.001)
+        >>> lr = env.get('learning_rate')
+        >>> print(f"Learning rate: {lr}")  # 0.001
+
+    With default value:
+
+    .. code-block:: python
+
+        >>> import brainstate.environ as env
+        >>>
+        >>> # Get with default
+        >>> batch_size = env.get('batch_size', default=32)
+        >>> print(f"Batch size: {batch_size}")  # 32
+
+    Context-aware retrieval:
+
+    .. code-block:: python
+
+        >>> import brainstate.environ as env
+        >>>
+        >>> env.set(temperature=1.0)
+        >>> print(env.get('temperature'))  # 1.0
+        >>>
+        >>> with env.context(temperature=0.5):
+        ...     print(env.get('temperature'))  # 0.5
+        >>>
+        >>> print(env.get('temperature'))  # 1.0
+
+    Notes
+    -----
+    Special keys 'platform' and 'host_device_count' are handled separately
+    and retrieve system-level information.
     """
-    if key == 'platform':
+    # Special cases for platform-specific parameters
+    if key == PLATFORM:
         return get_platform()
-
-    if key == 'host_device_count':
+    if key == HOST_DEVICE_COUNT:
         return get_host_device_count()
 
-    if key in DFAULT.contexts:
-        if len(DFAULT.contexts[key]) > 0:
-            return DFAULT.contexts[key][-1]
-    if key in DFAULT.settings:
-        return DFAULT.settings[key]
-
-    if default is _NOT_PROVIDE:
-        if desc is not None:
-            raise KeyError(
-                f"'{key}' is not found in the context. \n"
-                f"You can set it by `brainstate.share.context({key}=value)` "
-                f"locally or `brainstate.share.set({key}=value)` globally. \n"
-                f"Description: {desc}"
-            )
-        else:
-            raise KeyError(
-                f"'{key}' is not found in the context. \n"
-                f"You can set it by `brainstate.share.context({key}=value)` "
-                f"locally or `brainstate.share.set({key}=value)` globally."
-            )
+    # Check context stack first (most recent value)
+    with _ENV_STATE.locks[key]:
+        if key in _ENV_STATE.contexts and _ENV_STATE.contexts[key]:
+            return _ENV_STATE.contexts[key][-1]
+
+    # Check global settings
+    if key in _ENV_STATE.settings:
+        return _ENV_STATE.settings[key]
+
+    # Handle missing key
+    if default is _NOT_PROVIDED:
+        error_msg = f"Key '{key}' not found in environment."
+        if desc:
+            error_msg += f" Description: {desc}"
+        error_msg += (
+            f"\nSet it using:\n"
+            f"  - env.set({key}=value) for global setting\n"
+            f"  - env.context({key}=value) for temporary setting"
+        )
+        raise KeyError(error_msg)
+
     return default
 
 
-def all() -> dict:
+def all() -> Dict[str, Any]:
     """
-    Get all the current default computation environment.
+    Get all current environment settings.
+
+    This function returns a dictionary containing all active environment
+    settings, with context values taking precedence over global settings.
 
     Returns
     -------
-    r: dict
-      The current default computation environment.
+    dict
+        Dictionary of all current environment settings.
+
+    Examples
+    --------
+    .. code-block:: python
+
+        >>> import brainstate.environ as env
+        >>>
+        >>> # Set various parameters
+        >>> env.set(precision=32, dt=0.01, debug=True)
+        >>>
+        >>> # Get all settings
+        >>> settings = env.all()
+        >>> print(settings)
+        {'precision': 32, 'dt': 0.01, 'debug': True}
+
+        >>> # Context overrides
+        >>> with env.context(precision=64, new_param='test'):
+        ...     settings = env.all()
+        ...     print(settings['precision'])  # 64
+        ...     print(settings['new_param'])  # 'test'
+
+    Notes
+    -----
+    The returned dictionary is a snapshot and modifying it does not
+    affect the environment settings.
+    """
+    result = {}
+
+    # Add global settings
+    result.update(_ENV_STATE.settings)
+
+    # Override with context values (most recent)
+    for key, values in _ENV_STATE.contexts.items():
+        if values:
+            result[key] = values[-1]
+
+    return result
+
+
+def pop(key: str, default: Any = _NOT_PROVIDED) -> Any:
     """
-    r = dict()
-    for k, v in DFAULT.contexts.items():
-        if v:
-            r[k] = v[-1]
-    for k, v in DFAULT.settings.items():
-        if k not in r:
-            r[k] = v
-    return r
+    Remove and return a value from the global environment.
 
+    This function removes a key from the global environment settings and
+    returns its value. If the key is not found, it returns the default
+    value if provided, or raises KeyError.
 
-def get_dt():
-    """Get the numerical integrator precision.
+    Note that this function only affects global settings, not context values.
+    Keys in active contexts are not affected.
+
+    Parameters
+    ----------
+    key : str
+        The environment key to remove.
+    default : Any, optional
+        Default value to return if key is not found.
+        If not provided, raises KeyError for missing keys.
 
     Returns
     -------
-    dt : float
-        Numerical integration precision.
+    Any
+        The value that was removed from the environment.
+
+    Raises
+    ------
+    KeyError
+        If key is not found and no default is provided.
+    ValueError
+        If attempting to pop a key that is currently in a context.
+
+    Examples
+    --------
+    Basic usage:
+
+    .. code-block:: python
+
+        >>> import brainstate.environ as env
+        >>>
+        >>> # Set a value
+        >>> env.set(temp_param='temporary')
+        >>> print(env.get('temp_param'))  # 'temporary'
+        >>>
+        >>> # Pop the value
+        >>> value = env.pop('temp_param')
+        >>> print(value)  # 'temporary'
+        >>>
+        >>> # Value is now gone
+        >>> env.get('temp_param', default=None)  # None
+
+    With default value:
+
+    .. code-block:: python
+
+        >>> import brainstate.environ as env
+        >>>
+        >>> # Pop non-existent key with default
+        >>> value = env.pop('nonexistent', default='default_value')
+        >>> print(value)  # 'default_value'
+
+    Pop multiple values:
+
+    .. code-block:: python
+
+        >>> import brainstate.environ as env
+        >>>
+        >>> # Set multiple values
+        >>> env.set(param1='value1', param2='value2', param3='value3')
+        >>>
+        >>> # Pop them one by one
+        >>> v1 = env.pop('param1')
+        >>> v2 = env.pop('param2')
+        >>>
+        >>> # param3 still exists
+        >>> print(env.get('param3'))  # 'value3'
+
+    Context protection:
+
+    .. code-block:: python
+
+        >>> import brainstate.environ as env
+        >>>
+        >>> env.set(protected='global_value')
+        >>>
+        >>> with env.context(protected='context_value'):
+        ...     # Cannot pop a key that's in active context
+        ...     try:
+        ...         env.pop('protected')
+        ...     except ValueError as e:
+        ...         print("Cannot pop key in active context")
+
+    Notes
+    -----
+    - This function only removes keys from global settings
+    - Keys that are currently overridden in active contexts cannot be popped
+    - Special keys like 'platform' and 'host_device_count' can be popped but
+      their system-level values remain accessible through get_platform() etc.
+    - Registered callbacks are NOT triggered when popping values
     """
-    return get('dt')
+    # Check if key is currently in any active context
+    if key in _ENV_STATE.contexts and _ENV_STATE.contexts[key]:
+        raise ValueError(
+            f"Cannot pop key '{key}' while it is active in a context. "
+            f"The key is currently overridden in {len(_ENV_STATE.contexts[key])} context(s)."
+        )
+
+    # Check if key exists in global settings
+    if key in _ENV_STATE.settings:
+        # Remove and return the value
+        value = _ENV_STATE.settings.pop(key)
+
+        # Note: We don't trigger callbacks here as this is a removal operation
+        # If needed, users can register callbacks for removal separately
+
+        return value
+
+    # Key not found, handle default
+    if default is _NOT_PROVIDED:
+        raise KeyError(f"Key '{key}' not found in global environment settings.")
+
+    return default
+
 
+def set(
+    platform: Optional[PlatformType] = None,
+    host_device_count: Optional[int] = None,
+    precision: Optional[PrecisionType] = None,
+    dt: Optional[float] = None,
+    **kwargs
+) -> None:
+    """
+    Set global environment configuration.
 
-def get_mode() -> Mode:
-    """Get the default computing mode.
+    This function sets persistent global environment settings that remain
+    active until explicitly changed or the program terminates.
 
-    References
+    Parameters
     ----------
-    mode: Mode
-      The default computing mode.
+    platform : str, optional
+        Computing platform ('cpu', 'gpu', or 'tpu').
+    host_device_count : int, optional
+        Number of host devices for parallel computation.
+    precision : int or str, optional
+        Numerical precision (8, 16, 32, 64, or 'bf16').
+    mode : Mode, optional
+        Computation mode instance.
+    dt : float, optional
+        Time step for numerical integration.
+    **kwargs
+        Additional custom environment parameters.
+
+    Raises
+    ------
+    ValueError
+        If invalid platform or precision is specified.
+    TypeError
+        If mode is not a Mode instance.
+
+    Examples
+    --------
+    Basic configuration:
+
+    .. code-block:: python
+
+        >>> 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:
+
+    .. code-block:: python
+
+        >>> 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:
+
+    .. code-block:: python
+
+        >>> 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'
+
+    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
     """
-    return get('mode')
+    # Handle special parameters
+    if platform is not None:
+        set_platform(platform)
 
+    if host_device_count is not None:
+        set_host_device_count(host_device_count)
 
-def get_platform() -> str:
-    """Get the computing platform.
+    if precision is not None:
+        _validate_precision(precision)
+        _set_jax_precision(precision)
+        kwargs[PRECISION] = precision
 
-    Returns
-    -------
-    platform: str
-      Either 'cpu', 'gpu' or 'tpu'.
-    """
-    return devices()[0].platform
+    if dt is not None:
+        if not u.math.isscalar(dt):
+            raise TypeError(f"'{DT}' must be a scalar number, got {type(dt)}")
+        kwargs[DT] = dt
 
+    # Update global settings
+    _ENV_STATE.settings.update(kwargs)
 
-def get_host_device_count():
+    # Trigger registered callbacks
+    for key, value in kwargs.items():
+        if key in _ENV_STATE.functions:
+            try:
+                _ENV_STATE.functions[key](value)
+            except Exception as e:
+                warnings.warn(
+                    f"Callback for '{key}' raised an exception: {e}",
+                    RuntimeWarning
+                )
+
+
+def get_dt() -> float:
     """
-    Get the number of host devices.
+    Get the current numerical integration time step.
 
     Returns
     -------
-    n: int
-      The number of host devices.
+    float
+        The time step value.
+
+    Raises
+    ------
+    KeyError
+        If dt is not set.
+
+    Examples
+    --------
+    .. code-block:: python
+
+        >>> import brainstate.environ as env
+        >>>
+        >>> env.set(dt=0.01)
+        >>> dt = env.get_dt()
+        >>> print(f"Time step: {dt} ms")  # Time step: 0.01 ms
+        >>>
+        >>> # Use in computation
+        >>> with env.context(dt=0.001):
+        ...     fine_dt = env.get_dt()
+        ...     print(f"Fine time step: {fine_dt}")  # 0.001
     """
-    xla_flags = os.getenv("XLA_FLAGS", "")
-    match = re.search(r"--xla_force_host_platform_device_count=(\d+)", xla_flags)
-    return int(match.group(1)) if match else 1
+    return get(DT)
 
 
-def _get_precision() -> int | str:
+def get_platform() -> PlatformType:
     """
-    Get the default precision.
+    Get the current computing platform.
 
     Returns
     -------
-    precision: int
-      The default precision.
+    str
+        Platform name ('cpu', 'gpu', or 'tpu').
+
+    Examples
+    --------
+    .. code-block:: python
+
+        >>> import brainstate.environ as env
+        >>>
+        >>> platform = env.get_platform()
+        >>> print(f"Running on: {platform}")
+        >>>
+        >>> if platform == 'gpu':
+        ...     print("GPU acceleration available")
+        ... else:
+        ...     print(f"Using {platform.upper()}")
     """
-    return get('precision')
+    return devices()[0].platform
 
 
-def get_precision() -> int:
+def get_host_device_count() -> int:
     """
-    Get the default precision.
+    Get the number of host devices.
 
     Returns
     -------
-    precision: int
-      The default precision.
+    int
+        Number of host devices configured.
+
+    Examples
+    --------
+    .. code-block:: python
+
+        >>> import brainstate.environ as env
+        >>>
+        >>> # Get device count
+        >>> n_devices = env.get_host_device_count()
+        >>> print(f"Host devices: {n_devices}")
+        >>>
+        >>> # Configure for parallel computation
+        >>> if n_devices > 1:
+        ...     print(f"Can use {n_devices} devices for parallel computation")
     """
-    precision = get('precision')
-    if precision == 'bf16':
-        return 16
-    if isinstance(precision, int):
-        return precision
-    if isinstance(precision, str):
-        return int(precision)
-    raise ValueError(f'Unsupported precision: {precision}')
+    xla_flags = os.getenv("XLA_FLAGS", "")
+    match = re.search(r"--xla_force_host_platform_device_count=(\d+)", xla_flags)
+    return int(match.group(1)) if match else 1
 
 
-def set(
-    platform: str = None,
-    host_device_count: int = None,
-    precision: int | str = None,
-    mode: Mode = None,
-    **kwargs
-):
+def set_platform(platform: PlatformType) -> None:
     """
-    Set the global default computation environment.
-
+    Set the computing platform.
 
-
-    Args:
-      platform: str. The computing platform. Either 'cpu', 'gpu' or 'tpu'.
-      host_device_count: int. The number of host devices.
-      precision: int, str. The default precision.
-      mode: Mode. The computing mode.
-      **kwargs: dict. Other environment settings.
+    Parameters
+    ----------
+    platform : str
+        Platform to use ('cpu', 'gpu', or 'tpu').
+
+    Raises
+    ------
+    ValueError
+        If platform is not supported.
+
+    Examples
+    --------
+    .. code-block:: python
+
+        >>> import brainstate.environ as env
+        >>>
+        >>> # Set to GPU
+        >>> env.set_platform('gpu')
+        >>>
+        >>> # Verify platform
+        >>> print(env.get_platform())  # 'gpu'
+
+    Notes
+    -----
+    Platform changes only take effect at program start. Changing platform
+    after JAX initialization may not have the expected effect.
     """
-    if platform is not None:
-        set_platform(platform)
-    if host_device_count is not None:
-        set_host_device_count(host_device_count)
-    if precision is not None:
-        _set_jax_precision(precision)
-        kwargs['precision'] = precision
-    if mode is not None:
-        assert isinstance(mode, Mode), 'mode must be a Mode instance.'
-        kwargs['mode'] = mode
+    if platform not in SUPPORTED_PLATFORMS:
+        raise ValueError(
+            f"Platform must be one of {SUPPORTED_PLATFORMS}, got '{platform}'"
+        )
 
-    # set default environment
-    DFAULT.settings.update(kwargs)
+    config.update("jax_platform_name", platform)
 
-    # update the environment functions
-    for k, v in kwargs.items():
-        if k in DFAULT.functions:
-            DFAULT.functions[k](v)
+    # Trigger callbacks
+    if PLATFORM in _ENV_STATE.functions:
+        _ENV_STATE.functions[PLATFORM](platform)
 
 
-def set_host_device_count(n):
+def set_host_device_count(n: int) -> None:
     """
-    By default, XLA considers all CPU cores as one device. This utility tells XLA
-    that there are `n` host (CPU) devices available to use. As a consequence, this
-    allows parallel mapping in JAX :func:`jax.pmap` to work in CPU platform.
+    Set the number of host (CPU) devices.
 
-    .. note:: This utility only takes effect at the beginning of your program.
-        Under the hood, this sets the environment variable
-        `XLA_FLAGS=--xla_force_host_platform_device_count=[num_devices]`, where
-        `[num_device]` is the desired number of CPU devices `n`.
+    This function configures XLA to treat CPU cores as separate devices,
+    enabling parallel computation with jax.pmap on CPU.
 
-    .. warning:: Our understanding of the side effects of using the
-        `xla_force_host_platform_device_count` flag in XLA is incomplete. If you
-        observe some strange phenomenon when using this utility, please let us
-        know through our issue or forum page. More information is available in this
-        `JAX issue <https://github.com/google/jax/issues/1408>`_.
-
-    :param int n: number of devices to use.
+    Parameters
+    ----------
+    n : int
+        Number of host devices to configure.
+
+    Raises
+    ------
+    ValueError
+        If n is not a positive integer.
+
+    Examples
+    --------
+    .. code-block:: python
+
+        >>> import brainstate.environ as env
+        >>> import jax
+        >>>
+        >>> # Configure 4 CPU devices
+        >>> env.set_host_device_count(4)
+        >>>
+        >>> # Use with pmap
+        >>> def parallel_fn(x):
+        ...     return x * 2
+        >>>
+        >>> # This will work with 4 devices
+        >>> pmapped_fn = jax.pmap(parallel_fn)
+
+    Warnings
+    --------
+    This setting only takes effect at program start. The effects of using
+    xla_force_host_platform_device_count are not fully understood and may
+    cause unexpected behavior.
     """
+    if not isinstance(n, int) or n < 1:
+        raise ValueError(f"Host device count must be a positive integer, got {n}")
+
+    # Update XLA flags
     xla_flags = os.getenv("XLA_FLAGS", "")
-    xla_flags = re.sub(r"--xla_force_host_platform_device_count=\S+", "", xla_flags).split()
-    os.environ["XLA_FLAGS"] = " ".join(["--xla_force_host_platform_device_count={}".format(n)] + xla_flags)
+    xla_flags = re.sub(
+        r"--xla_force_host_platform_device_count=\S+",
+        "",
+        xla_flags
+    ).split()
 
-    # update the environment functions
-    if 'host_device_count' in DFAULT.functions:
-        DFAULT.functions['host_device_count'](n)
+    os.environ["XLA_FLAGS"] = " ".join(
+        [f"--xla_force_host_platform_device_count={n}"] + xla_flags
+    )
 
+    # Trigger callbacks
+    if HOST_DEVICE_COUNT in _ENV_STATE.functions:
+        _ENV_STATE.functions[HOST_DEVICE_COUNT](n)
 
-def set_platform(platform: str):
-    """
-    Changes platform to CPU, GPU, or TPU. This utility only takes
-    effect at the beginning of your program.
 
-    Args:
-      platform: str. The computing platform. Either 'cpu', 'gpu' or 'tpu'.
+def set_precision(precision: PrecisionType) -> None:
+    """
+    Set the global numerical precision.
 
-    Raises:
-      ValueError: If the platform is not in ['cpu', 'gpu', 'tpu'].
+    Parameters
+    ----------
+    precision : int or str
+        Precision to use (8, 16, 32, 64, or 'bf16').
+
+    Raises
+    ------
+    ValueError
+        If precision is not supported.
+
+    Examples
+    --------
+    .. code-block:: python
+
+        >>> import brainstate.environ as env
+        >>> import jax.numpy as jnp
+        >>>
+        >>> # Set to 64-bit precision
+        >>> env.set_precision(64)
+        >>>
+        >>> # Arrays will use float64 by default
+        >>> x = jnp.array([1.0, 2.0, 3.0])
+        >>> print(x.dtype)  # float64
+        >>>
+        >>> # Set to bfloat16 for efficiency
+        >>> env.set_precision('bf16')
     """
-    assert platform in ['cpu', 'gpu', 'tpu']
-    config.update("jax_platform_name", platform)
+    _validate_precision(precision)
+    _set_jax_precision(precision)
+    _ENV_STATE.settings[PRECISION] = precision
 
-    # update the environment functions
-    if 'platform' in DFAULT.functions:
-        DFAULT.functions['platform'](platform)
+    # Trigger callbacks
+    if PRECISION in _ENV_STATE.functions:
+        _ENV_STATE.functions[PRECISION](precision)
 
 
-def _set_jax_precision(precision: int | str):
+def get_precision() -> int:
     """
-    Set the default precision.
+    Get the current numerical precision as an integer.
 
-    Args:
-      precision: int. The default precision.
+    Returns
+    -------
+    int
+        Precision in bits (8, 16, 32, or 64).
+
+    Examples
+    --------
+    .. code-block:: python
+
+        >>> import brainstate.environ as env
+        >>>
+        >>> env.set_precision(32)
+        >>> bits = env.get_precision()
+        >>> print(f"Using {bits}-bit precision")  # Using 32-bit precision
+        >>>
+        >>> # Special handling for bfloat16
+        >>> env.set_precision('bf16')
+        >>> print(env.get_precision())  # 16
+
+    Notes
+    -----
+    'bf16' (bfloat16) is reported as 16-bit precision.
     """
-    # assert precision in [64, 32, 16, 'bf16', 8], f'Precision must be in [64, 32, 16, "bf16", 8]. But got {precision}.'
-    if precision in [64, '64']:
+    precision = get(PRECISION, default=DEFAULT_PRECISION)
+
+    if precision == 'bf16':
+        return 16
+    elif isinstance(precision, str):
+        return int(precision)
+    elif isinstance(precision, int):
+        return precision
+    else:
+        raise ValueError(f"Invalid precision type: {type(precision)}")
+
+
+def _validate_precision(precision: PrecisionType) -> None:
+    """Validate precision value."""
+    if precision not in SUPPORTED_PRECISIONS and str(precision) not in map(str, SUPPORTED_PRECISIONS):
+        raise ValueError(
+            f"Precision must be one of {SUPPORTED_PRECISIONS}, got {precision}"
+        )
+
+
+def _get_precision() -> PrecisionType:
+    """Get raw precision value (including 'bf16')."""
+    return get(PRECISION, default=DEFAULT_PRECISION)
+
+
+def _set_jax_precision(precision: PrecisionType) -> None:
+    """Configure JAX precision settings."""
+    # Enable/disable 64-bit mode
+    if precision in (64, '64'):
         config.update("jax_enable_x64", True)
     else:
         config.update("jax_enable_x64", False)
 
 
-@functools.lru_cache()
-def _get_uint(precision: int):
-    if precision in [64, '64']:
+@functools.lru_cache(maxsize=16)
+def _get_uint(precision: PrecisionType) -> DTypeLike:
+    """Get unsigned integer type for given precision."""
+    if precision in (64, '64'):
         return np.uint64
-    elif precision in [32, '32']:
+    elif precision in (32, '32'):
         return np.uint32
-    elif precision in [16, '16', 'bf16']:
+    elif precision in (16, '16', 'bf16'):
         return np.uint16
-    elif precision in [8, '8']:
+    elif precision in (8, '8'):
         return np.uint8
     else:
-        raise ValueError(f'Unsupported precision: {precision}')
+        raise ValueError(f"Unsupported precision: {precision}")
 
 
-@functools.lru_cache()
-def _get_int(precision: int):
-    if precision in [64, '64']:
+@functools.lru_cache(maxsize=16)
+def _get_int(precision: PrecisionType) -> DTypeLike:
+    """Get integer type for given precision."""
+    if precision in (64, '64'):
         return np.int64
-    elif precision in [32, '32']:
+    elif precision in (32, '32'):
         return np.int32
-    elif precision in [16, '16', 'bf16']:
+    elif precision in (16, '16', 'bf16'):
         return np.int16
-    elif precision in [8, '8']:
+    elif precision in (8, '8'):
         return np.int8
     else:
-        raise ValueError(f'Unsupported precision: {precision}')
+        raise ValueError(f"Unsupported precision: {precision}")
 
 
-@functools.lru_cache()
-def _get_float(precision: int):
-    if precision in [64, '64']:
+@functools.lru_cache(maxsize=16)
+def _get_float(precision: PrecisionType) -> DTypeLike:
+    """Get floating-point type for given precision."""
+    if precision in (64, '64'):
         return np.float64
-    elif precision in [32, '32']:
+    elif precision in (32, '32'):
         return np.float32
-    elif precision in [16, '16']:
+    elif precision in (16, '16'):
         return np.float16
-    elif precision in ['bf16']:
+    elif precision == 'bf16':
         return jnp.bfloat16
-    elif precision in [8, '8']:
+    elif precision in (8, '8'):
         return jnp.float8_e5m2
     else:
-        raise ValueError(f'Unsupported precision: {precision}')
+        raise ValueError(f"Unsupported precision: {precision}")
 
 
-@functools.lru_cache()
-def _get_complex(precision: int):
-    if precision == [64, '64']:
+@functools.lru_cache(maxsize=16)
+def _get_complex(precision: PrecisionType) -> DTypeLike:
+    """Get complex type for given precision."""
+    if precision in (64, '64'):
         return np.complex128
-    elif precision == [32, '32']:
-        return np.complex64
-    elif precision in [16, '16', 'bf16']:
-        return np.complex64
-    elif precision == [8, '8']:
+    elif precision in (32, '32', 16, '16', 'bf16', 8, '8'):
         return np.complex64
     else:
-        raise ValueError(f'Unsupported precision: {precision}')
+        raise ValueError(f"Unsupported precision: {precision}")
 
 
 def dftype() -> DTypeLike:
     """
-    Default floating data type.
+    Get the default floating-point data type.
 
-    This function returns the default floating data type based on the current precision.
-    If you want the data dtype is changed with the setting of the precision by ``brainstate.environ.set(precision)``,
-    you can use this function to get the default floating data type, and create the data by using ``dtype=dftype()``.
-
-    For example, if the precision is set to 32, the default floating data type is ``np.float32``.
-
-    >>> import brainstate as brainstate
-    >>> import numpy as np
-    >>> with brainstate.environ.context(precision=32):
-    ...    a = np.zeros(1, dtype=brainstate.environ.dftype())
-    >>> print(a.dtype)
+    This function returns the appropriate floating-point type based on
+    the current precision setting, allowing dynamic type selection.
 
     Returns
     -------
-    float_dtype: DTypeLike
-      The default floating data type.
+    DTypeLike
+        Default floating-point data type.
+
+    Examples
+    --------
+    .. code-block:: python
+
+        >>> import brainstate.environ as env
+        >>> import jax.numpy as jnp
+        >>>
+        >>> # With 32-bit precision
+        >>> env.set(precision=32)
+        >>> x = jnp.zeros(10, dtype=env.dftype())
+        >>> print(x.dtype)  # float32
+        >>>
+        >>> # With 64-bit precision
+        >>> with env.context(precision=64):
+        ...     y = jnp.ones(5, dtype=env.dftype())
+        ...     print(y.dtype)  # float64
+        >>>
+        >>> # With bfloat16
+        >>> env.set(precision='bf16')
+        >>> z = jnp.array([1, 2, 3], dtype=env.dftype())
+        >>> print(z.dtype)  # bfloat16
+
+    See Also
+    --------
+    ditype : Default integer type
+    dutype : Default unsigned integer type
+    dctype : Default complex type
     """
     return _get_float(_get_precision())
 
 
 def ditype() -> DTypeLike:
     """
-    Default integer data type.
-
-    This function returns the default integer data type based on the current precision.
-    If you want the data dtype is changed with the setting of the precision by ``brainstate.environ.set(precision)``,
-    you can use this function to get the default integer data type, and create the data by using ``dtype=ditype()``.
-
-    For example, if the precision is set to 32, the default integer data type is ``np.int32``.
+    Get the default integer data type.
 
-    >>> import brainstate as brainstate
-    >>> import numpy as np
-    >>> with brainstate.environ.context(precision=32):
-    ...    a = np.zeros(1, dtype=brainstate.environ.ditype())
-    >>> print(a.dtype)
-    int32
+    This function returns the appropriate integer type based on
+    the current precision setting.
 
     Returns
     -------
-    int_dtype: DTypeLike
-      The default integer data type.
+    DTypeLike
+        Default integer data type.
+
+    Examples
+    --------
+    .. code-block:: python
+
+        >>> import brainstate.environ as env
+        >>> import jax.numpy as jnp
+        >>>
+        >>> # With 32-bit precision
+        >>> env.set(precision=32)
+        >>> indices = jnp.arange(10, dtype=env.ditype())
+        >>> print(indices.dtype)  # int32
+        >>>
+        >>> # With 64-bit precision
+        >>> with env.context(precision=64):
+        ...     big_indices = jnp.arange(1000, dtype=env.ditype())
+        ...     print(big_indices.dtype)  # int64
+
+    See Also
+    --------
+    dftype : Default floating-point type
+    dutype : Default unsigned integer type
     """
     return _get_int(_get_precision())
 
 
 def dutype() -> DTypeLike:
     """
-    Default unsigned integer data type.
+    Get the default unsigned integer data type.
 
-    This function returns the default unsigned integer data type based on the current precision.
-    If you want the data dtype is changed with the setting of the precision
-    by ``brainstate.environ.set(precision)``, you can use this function to get the default
-    unsigned integer data type, and create the data by using ``dtype=dutype()``.
-
-    For example, if the precision is set to 32, the default unsigned integer data type is ``np.uint32``.
-
-    >>> import brainstate as brainstate
-    >>> import numpy as np
-    >>> with brainstate.environ.context(precision=32):
-    ...   a = np.zeros(1, dtype=brainstate.environ.dutype())
-    >>> print(a.dtype)
-    uint32
+    This function returns the appropriate unsigned integer type based on
+    the current precision setting.
 
     Returns
     -------
-    uint_dtype: DTypeLike
-      The default unsigned integer data type.
+    DTypeLike
+        Default unsigned integer data type.
+
+    Examples
+    --------
+    .. code-block:: python
+
+        >>> import brainstate.environ as env
+        >>> import jax.numpy as jnp
+        >>>
+        >>> # With 32-bit precision
+        >>> env.set(precision=32)
+        >>> counts = jnp.array([10, 20, 30], dtype=env.dutype())
+        >>> print(counts.dtype)  # uint32
+        >>>
+        >>> # With 16-bit precision
+        >>> with env.context(precision=16):
+        ...     small_counts = jnp.array([1, 2, 3], dtype=env.dutype())
+        ...     print(small_counts.dtype)  # uint16
+
+    See Also
+    --------
+    ditype : Default signed integer type
     """
     return _get_uint(_get_precision())
 
 
 def dctype() -> DTypeLike:
     """
-    Default complex data type.
-
-    This function returns the default complex data type based on the current precision.
-    If you want the data dtype is changed with the setting of the precision by ``brainstate.environ.set(precision)``,
-    you can use this function to get the default complex data type, and create the data by using ``dtype=dctype()``.
+    Get the default complex data type.
 
-    For example, if the precision is set to 32, the default complex data type is ``np.complex64``.
-
-    >>> import brainstate as brainstate
-    >>> import numpy as np
-    >>> with brainstate.environ.context(precision=32):
-    ...    a = np.zeros(1, dtype=brainstate.environ.dctype())
-    >>> print(a.dtype)
-    complex64
+    This function returns the appropriate complex type based on
+    the current precision setting.
 
     Returns
     -------
-    complex_dtype: DTypeLike
-      The default complex data type.
+    DTypeLike
+        Default complex data type.
+
+    Examples
+    --------
+    .. code-block:: python
+
+        >>> import brainstate.environ as env
+        >>> import jax.numpy as jnp
+        >>>
+        >>> # With 32-bit precision
+        >>> env.set(precision=32)
+        >>> z = jnp.array([1+2j, 3+4j], dtype=env.dctype())
+        >>> print(z.dtype)  # complex64
+        >>>
+        >>> # With 64-bit precision
+        >>> with env.context(precision=64):
+        ...     w = jnp.array([5+6j], dtype=env.dctype())
+        ...     print(w.dtype)  # complex128
+
+    Notes
+    -----
+    Complex128 is only available with 64-bit precision.
+    All other precisions use complex64.
     """
     return _get_complex(_get_precision())
 
 
-def tolerance():
-    if get_precision() == 64:
+def tolerance() -> jnp.ndarray:
+    """
+    Get numerical tolerance based on current precision.
+
+    This function returns an appropriate tolerance value for numerical
+    comparisons based on the current precision setting.
+
+    Returns
+    -------
+    jnp.ndarray
+        Tolerance value as a scalar array.
+
+    Examples
+    --------
+    .. code-block:: python
+
+        >>> import brainstate.environ as env
+        >>> import jax.numpy as jnp
+        >>>
+        >>> # Different tolerances for different precisions
+        >>> env.set(precision=64)
+        >>> tol64 = env.tolerance()
+        >>> print(f"64-bit tolerance: {tol64}")  # 1e-12
+        >>>
+        >>> env.set(precision=32)
+        >>> tol32 = env.tolerance()
+        >>> print(f"32-bit tolerance: {tol32}")  # 1e-5
+        >>>
+        >>> # Use in numerical comparisons
+        >>> def are_close(a, b):
+        ...     return jnp.abs(a - b) < env.tolerance()
+
+    Notes
+    -----
+    Tolerance values:
+    - 64-bit: 1e-12
+    - 32-bit: 1e-5
+    - 16-bit and below: 1e-2
+    """
+    precision = get_precision()
+
+    if precision == 64:
         return jnp.array(1e-12, dtype=np.float64)
-    elif get_precision() == 32:
+    elif precision == 32:
         return jnp.array(1e-5, dtype=np.float32)
     else:
         return jnp.array(1e-2, dtype=np.float16)
 
 
-def register_default_behavior(key: str, behavior: Callable, replace_if_exist: bool = False):
+def register_default_behavior(
+    key: str,
+    behavior: Callable[[Any], None],
+    replace_if_exist: bool = False
+) -> None:
     """
-    Register a default behavior for a specific global key parameter.
+    Register a callback for environment parameter changes.
 
-    For example, you can register a default behavior for the key 'dt' by::
+    This function allows you to register custom behaviors that are
+    triggered whenever a specific environment parameter is modified.
 
-    >>> import brainstate as brainstate
-    >>> def dt_behavior(dt):
-    ...     print(f'Set the default dt to {dt}.')
-    ...
-    >>> brainstate.environ.register_default_behavior('dt', dt_behavior)
+    Parameters
+    ----------
+    key : str
+        Environment parameter key to monitor.
+    behavior : Callable[[Any], None]
+        Callback function that receives the new value.
+    replace_if_exist : bool, default=False
+        Whether to replace existing callback for this key.
+
+    Raises
+    ------
+    TypeError
+        If behavior is not callable.
+    ValueError
+        If key already has a registered behavior and replace_if_exist is False.
+
+    Examples
+    --------
+    Basic callback registration:
+
+    .. code-block:: python
+
+        >>> import brainstate.environ as env
+        >>>
+        >>> # Define a callback
+        >>> def on_dt_change(new_dt):
+        ...     print(f"Time step changed to: {new_dt}")
+        >>>
+        >>> # Register the callback
+        >>> env.register_default_behavior('dt', on_dt_change)
+        >>>
+        >>> # Callback is triggered on changes
+        >>> env.set(dt=0.01)  # Prints: Time step changed to: 0.01
+        >>>
+        >>> with env.context(dt=0.001):  # Prints: Time step changed to: 0.001
+        ...     pass  # Prints: Time step changed to: 0.01 (on exit)
+
+    Complex behavior with validation:
+
+    .. code-block:: python
+
+        >>> import brainstate.environ as env
+        >>>
+        >>> def validate_batch_size(size):
+        ...     if not isinstance(size, int) or size <= 0:
+        ...         raise ValueError(f"Invalid batch size: {size}")
+        ...     if size > 1024:
+        ...         print(f"Warning: Large batch size {size} may cause OOM")
+        >>>
+        >>> env.register_default_behavior('batch_size', validate_batch_size)
+        >>>
+        >>> # Valid setting
+        >>> env.set(batch_size=32)  # OK
+        >>>
+        >>> # Invalid setting
+        >>> # env.set(batch_size=-1)  # Raises ValueError
+
+    Replacing existing behavior:
+
+    .. code-block:: python
+
+        >>> import brainstate.environ as env
+        >>>
+        >>> def old_behavior(value):
+        ...     print(f"Old: {value}")
+        >>>
+        >>> def new_behavior(value):
+        ...     print(f"New: {value}")
+        >>>
+        >>> env.register_default_behavior('key', old_behavior)
+        >>> env.register_default_behavior('key', new_behavior, replace_if_exist=True)
+        >>>
+        >>> env.set(key='test')  # Prints: New: test
+
+    See Also
+    --------
+    unregister_default_behavior : Remove registered callbacks
+    list_registered_behaviors : List all registered callbacks
+    """
+    if not isinstance(key, str):
+        raise TypeError(f"Key must be a string, got {type(key)}")
 
-    Then, when you set the default dt by `brainstate.environ.set(dt=0.1)`, the behavior
-    `dt_behavior` will be called with
-    `dt_behavior(0.1)`.
+    if not callable(behavior):
+        raise TypeError(f"Behavior must be callable, got {type(behavior)}")
 
-    >>> brainstate.environ.set(dt=0.1)
-    Set the default dt to 0.1.
-    >>> with brainstate.environ.context(dt=0.2):
-    ...     pass
-    Set the default dt to 0.2.
-    Set the default dt to 0.1.
+    if key in _ENV_STATE.functions and not replace_if_exist:
+        raise ValueError(
+            f"Behavior for key '{key}' already registered. "
+            f"Use replace_if_exist=True to override."
+        )
 
+    _ENV_STATE.functions[key] = behavior
 
-    Args:
-      key: str. The key to register.
-      behavior: Callable. The behavior to register. It should be a callable.
-      replace_if_exist: bool. Whether to replace the behavior if the key has been registered.
 
+def unregister_default_behavior(key: str) -> bool:
+    """
+    Remove a registered callback for an environment parameter.
+
+    Parameters
+    ----------
+    key : str
+        Environment parameter key.
+
+    Returns
+    -------
+    bool
+        True if a callback was removed, False if no callback existed.
+
+    Examples
+    --------
+    .. code-block:: python
+
+        >>> import brainstate.environ as env
+        >>>
+        >>> # Register a callback
+        >>> def callback(value):
+        ...     print(f"Value: {value}")
+        >>>
+        >>> env.register_default_behavior('param', callback)
+        >>>
+        >>> # Remove the callback
+        >>> removed = env.unregister_default_behavior('param')
+        >>> print(f"Callback removed: {removed}")  # True
+        >>>
+        >>> # No callback triggers now
+        >>> env.set(param='test')  # No output
+        >>>
+        >>> # Removing non-existent callback
+        >>> removed = env.unregister_default_behavior('nonexistent')
+        >>> print(f"Callback removed: {removed}")  # False
+    """
+    if key in _ENV_STATE.functions:
+        del _ENV_STATE.functions[key]
+        return True
+    return False
+
+
+def list_registered_behaviors() -> List[str]:
+    """
+    List all keys with registered callbacks.
+
+    Returns
+    -------
+    list of str
+        Keys that have registered behavior callbacks.
+
+    Examples
+    --------
+    .. code-block:: python
+
+        >>> import brainstate.environ as env
+        >>>
+        >>> # Register some callbacks
+        >>> env.register_default_behavior('param1', lambda x: None)
+        >>> env.register_default_behavior('param2', lambda x: None)
+        >>>
+        >>> # List registered behaviors
+        >>> behaviors = env.list_registered_behaviors()
+        >>> print(f"Registered: {behaviors}")  # ['param1', 'param2']
+        >>>
+        >>> # Check if specific behavior is registered
+        >>> if 'dt' in behaviors:
+        ...     print("dt has a registered callback")
     """
-    assert isinstance(key, str), 'key must be a string.'
-    assert callable(behavior), 'behavior must be a callable.'
-    if not replace_if_exist:
-        assert key not in DFAULT.functions, f'{key} has been registered.'
-    DFAULT.functions[key] = behavior
+    return list(_ENV_STATE.functions.keys())
 
 
-set(precision=32)
+# Initialize default precision on module load
+set(precision=DEFAULT_PRECISION)