brainstate module

Core State Classes

State classes are the fundamental building blocks for managing dynamic data in BrainState. They provide a unified interface for tracking, tracing, and transforming stateful computations.

Basic State Types

Basic state types provide semantic distinctions for different data lifecycles in your program.

State	A generic class representing a dynamic data pointer in the BrainState framework.
ShortTermState	A class representing short-term state in a program.
LongTermState	The long-term state, which is used to store the long-term data in the program.
ParamState	The parameter state, which is used to store the trainable parameters in the model.
BatchState	The batch state, which is used to store the batch data in the program.
DelayState	Short-term state for storing delay data.

Hidden State Types

Hidden state types are designed for recurrent neural networks and eligibility trace-based learning, with special support for BrainScale online learning integration.

• HiddenState: Single hidden state variable for neurons or synapses (equivalent to brainstate.HiddenState)

• HiddenGroupState: Multiple hidden states stored in the last dimension of a single array

• HiddenTreeState: Multiple hidden states with different units, stored as a PyTree structure

HiddenState	Represents hidden state variables in neurons or synapses.
HiddenGroupState	A group of multiple hidden states for eligibility trace-based learning.
HiddenTreeState	A pytree of multiple hidden states for eligibility trace-based learning.

Special State Types

Special-purpose state types for advanced use cases and PyTree integration.

FakeState	The faked state, which is used to store the faked data in the program.
TreefyState	The state as a pytree.

State Management

Tools for managing collections of states and tracking state access patterns during program execution.

State Collections

Organize and manipulate multiple states as cohesive units.

StateDictManager

State stack, for collecting all State used in the program.

State Tracing

Track state read/write operations for automatic differentiation and program transformation.

StateTraceStack

A stack for tracing and managing states during program execution.

State Utilities

Helper functions and context managers for working with states effectively.

Context Managers

Control state behavior within specific code blocks.

check_state_value_tree	The contex manager to check weather the tree structure of the state value keeps consistently.
check_state_jax_tracer	The context manager to check whether the state is valid to trace.
catch_new_states	A context manager that catches and tracks new states created within its scope.

Helper Functions

Utility functions for common state operations.

maybe_state

Extracts the value from a State object if given, otherwise returns the input value.

Error Handling

Custom exceptions for state-related errors and debugging.

Exception Classes

BrainStateError	A custom exception class for BrainState-related errors.
BatchAxisError	Exception raised for errors related to batch axis operations.
TraceContextError	Exception raised for trace-context violations.

State Hooks

The hook system observes and intercepts state operations — reads, writes, initialization, and restoration. Hooks can be registered globally and are managed per-process, enabling logging, validation, and instrumentation of stateful computations without modifying the states themselves.

Hook Contexts

Context objects passed to hooks describing the operation being performed.

HookContext	Base context object passed to all hooks.
ReadHookContext	Context for read hooks.
WriteHookContext	Base context for write operations.
MutableWriteHookContext	Context for write_before hooks with transformation capabilities.
RestoreHookContext	Context for restore hooks.
InitHookContext	Context for initialization hooks.

Hook Core

Core hook abstractions: the Hook base class and the HookHandle returned on registration.

Hook	Base hook class for all hook types.
HookHandle	Handle for managing a registered hook.

Hook Manager and Registry

Manage hook lifecycles and global registration.

HookManager	Manager for hooks on a single State instance.
HookConfig	Configuration for hook error handling and behavior.
GlobalHookRegistry	Singleton registry for global hooks that apply to all State instances.

Global Hook Functions

Convenience functions for registering and querying globally-installed hooks.

register_state_hook	Register a global hook that applies to all State instances.
unregister_state_hook	Unregister a global hook using its handle.
clear_state_hooks	Clear all global hooks, optionally filtered by type.
has_state_hooks	Check if any global hooks are registered.
list_state_hooks	List all registered global hooks, optionally filtered by type.

Hook Exceptions

Exceptions and warnings raised by the hook system.

HookError	Base exception for hook-related errors.
HookExecutionError	Exception raised when a hook execution fails.
HookRegistrationError	Exception raised when hook registration fails.
HookCancellationError	Exception raised when a hook cancels an operation.
HookWarning	Warning for hook-related issues.