register_default_behavior

register_default_behavior#

class brainstate.environ.register_default_behavior(key, behavior, replace_if_exist=False, *, env=None)[source]#

Register a callback for environment parameter changes.

This function allows you to register custom behaviors that are triggered whenever a specific environment parameter is modified.

Parameters:

key (str) – Environment parameter key to monitor.
behavior (Callable[[Any], None]) – Callback function that receives the new value.
replace_if_exist (bool) – Whether to replace existing callback for this key.
env (EnvironmentState | None) – The environment state to modify. If None, uses the global environment.

Raises:

TypeError – If behavior is not callable.
ValueError – If key already has a registered behavior and replace_if_exist is False.

Return type:

None

Examples

Basic callback registration:

>>> 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:

>>> 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:

>>> 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

Using custom environment:

>>> import brainstate.environ as env
>>>
>>> custom_env = env.EnvironmentState()
>>> env.register_default_behavior('param', lambda x: print(f"Value: {x}"), env=custom_env)
>>> env.set(param='test', env=custom_env)  # Prints: Value: test

register_default_behavior

Contents

register_default_behavior#