brainstate.register_state_hook

brainstate.register_state_hook#

brainstate.register_state_hook(hook_type, callback, priority=0, name=None, enabled=True)[source]#

Register a global hook that applies to all State instances.

Global hooks execute before instance-specific hooks.

Parameters:

  • hook_type (Literal['read', 'write_before', 'write_after', 'restore', 'init']) – Type of hook (‘read’, ‘write_before’, ‘write_after’, ‘restore’, ‘init’)

  • callback (Callable[[HookContext], Any]) – Callable that receives a HookContext

  • priority (int) – Execution priority (higher = earlier, default 0)

  • name (str | None) – Optional name for the hook

  • enabled (bool) – Whether the hook is enabled initially (default True)

Return type:

HookHandle

Examples

>>> import brainstate
>>> def validate_all_writes(ctx):
...     if hasattr(ctx.value, 'shape'):
...         print(f"Writing array with shape {ctx.value.shape}")
>>> handle = brainstate.register_state_hook('write_before', validate_all_writes)