brainstate.util module

Dict Operation

flat_mapping	Flatten a nested mapping into a flat mapping with tuple or string keys.
nest_mapping	Unflatten a mapping by converting tuple/string keys back to nested structure.
NestedDict	A pytree-like nested mapping structure for organizing hierarchical data.
FlattedDict	A pytree-like structure that contains a Mapping from strings or integers to leaves.

Functions for flattening, merging, and freezing nested dictionaries.

flatten_dict	Flatten a nested dictionary.
unflatten_dict	Unflatten a dictionary with separated keys.
merge_dicts	Merge multiple dictionaries.
freeze	Convert a mapping to a FrozenDict.
unfreeze	Convert a FrozenDict to a regular dict.
copy
pop
is_dataclass	Check whether a class was created by the brainstate dataclass() decorator.

Filter Operation

to_predicate converts a filter specification into a predicate; the remaining classes are composable filter primitives for selecting states and tree leaves.

to_predicate

Convert a Filter to a predicate function.

All	A filter class that combines multiple filters using a logical AND operation.
Any	Combine multiple filters using logical OR operation.
Not	A filter class that negates the result of another filter.
Nothing	Filter that matches no objects.
Everything	Filter that matches all objects.
OfType	Filter objects based on their type.
WithTag	Filter objects that have a specific tag attribute.
PathContains	Filter objects based on whether their path contains a specific key.

Pretty Representation

PrettyType	Configuration for pretty representation of objects.
PrettyAttr	Configuration for pretty representation of attributes.
PrettyRepr	Interface for pretty representation of objects.
PrettyMapping	Pretty representation of a mapping.
MappingReprMixin	Mapping mixin for pretty representation.
PrettyDict	Base dictionary class with pretty representation and tree utilities.
PrettyList	A list subclass with pretty representation and JAX pytree compatibility.
PrettyObject	Base class for generating pretty representations of tree-like structures.

pretty_repr	Create a pretty string representation.
yield_unique_pretty_repr_items	Generate a pretty representation of an object while avoiding duplicate representations.

Struct Operation

dataclass	Create a dataclass that works with JAX transformations.
field	Create a dataclass field with JAX pytree metadata.
PyTreeNode	Base class for creating JAX-compatible pytree nodes.
FrozenDict	An immutable dictionary that works as a JAX pytree.

Other Operations

clear_buffer_memory	Clear on-device memory buffers and optionally compilation cache.
not_instance_eval	Create a partial function to check if input is NOT an instance of given classes.
is_instance_eval	Create a partial function to check if input IS an instance of given classes.
DictManager	Enhanced dictionary for managing collections in BrainState.
DotDict	Dictionary with dot notation access to nested keys.
StateJaxTracer	Snapshot of the active JAX trace used to detect cross-trace State leakage.
BoundedCache	A thread-safe LRU cache with bounded size.

get_unique_name	Get a unique name for the given object type.
split_total	Calculate the number of epochs for simulation based on a total and a fraction.