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.

Filter Operation

to_predicate

Convert a Filter to a predicate function.

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.

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.