brainstate.graph moduel

Most of these APIs are adapted from Flax (google/flax). It enables the structure-preserving State retrieval and manipulatio in the brainstate.

A clean-room, flat-IR engine for flattening object graphs of Nodes and States into a static structure plus a dynamic state mapping, and back.

The engine is layered in three tiers:

• Flat IR (_graphdef) — a GraphDef is a root Edge plus a flat table of NodeSpecs. Graph nodes are hoisted into one index-keyed table (so sharing and cycles are encoded by integer index), while pytree containers are embedded inline. The whole structure hashes in a single cached pass and is registered as a JAX static pytree.

• Traversal kernel (_walk) — one depth-first primitive backs iter_leaf() / iter_node() (and states / nodes), plus the node-type registry behind register_graph_node_type().

• Encode / decode (_flatten) — flatten() encodes in one pre-order pass; unflatten() decodes in three linear passes (materialize states, create node shells, fill), so reconstruction over graph nodes never recurses and cycles/sharing resolve regardless of fill order.

Operations (treefy_split() / treefy_merge() / states() / nodes() / …) and graph/pytree conversion (graph_to_tree() / tree_to_graph()) are thin layers on top.

Graph Node

Node	Base class for all mutable graph nodes in the BrainState framework.

Graph Operation

pop_states	Remove and return ``State``s matching the filters (deduped by identity).
nodes	Collect graph nodes as ``FlattedDict``(s), optionally filtered.
states	Collect State objects from a graph node as ``FlattedDict``(s).
treefy_states	Return the treefy state mapping of node, optionally filtered.
update_states	Update node in place from one or more state mappings.
flatten	Flatten a graph node into a (GraphDef, NestedDict) pair.
unflatten	Reconstruct a graph node from a GraphDef and a state mapping.
treefy_split	Split a graph node into a GraphDef and one or more state mappings.
treefy_merge	Reconstruct a node from a GraphDef and one or more state mappings.
iter_leaf	Iterate (path, value) over every leaf in the graph.
iter_node	Iterate (path, graph_node) over every graph node in the graph.
clone	Deep-copy node via split/merge (shared references preserved).
graphdef	Return the GraphDef of node.

Context Management

Context managers for handling complex state updates during graph transformations. These utilities enable splitting and merging graph states in a thread-safe manner.

split_context	Context manager for splitting multiple graph nodes sharing a reference index.
merge_context	Context manager for merging multiple graph nodes sharing a reference index.

Graph Conversion

Utilities for converting between graph and tree representations, enabling flexible manipulation of nested module structures.

graph_to_tree	Convert a pytree that may contain graph nodes into a pure pytree.
tree_to_graph	Convert a pytree of NodeStates back into graph nodes.
NodeStates	A JAX pytree wrapper that carries both a GraphDef and one or more state mappings.

Graph Definition Classes

Core classes for representing graph structure, node definitions, and references. These classes provide the foundation for graph operations and state management.

GraphDef	The static structure of an object graph.
NodeSpec	The static record for a single graph node (never a pytree).
NodeEdge	A reference to a graph node by its global index.
StateEdge	A deduplicated State leaf, referenced by global index.
StateLeafEdge	A bare TreefyState leaf (carried directly, never deduplicated).
PytreeEdge	An inline pytree container.
StaticEdge	An inline, hashable static value carried directly in the structure.
Static	An empty pytree node that treats its inner value as static.
RefMap	A mutable mapping that keys entries by object identity (id).
register_graph_node_type	Register a custom mutable graph-node type with the graph engine.