flatten

Contents

flatten#

class brainstate.graph.flatten(node, /, ref_index=None, treefy_state=True)#

Flatten a graph node into a (GraphDef, NestedDict) pair.

Graph nodes and State objects are hoisted into a single global index space (deduplicated by identity, so shared references and cycles collapse to one index); pytree containers are embedded inline and re-expanded per occurrence.

Parameters:

  • node (Any) – The root graph node or pytree to flatten.

  • ref_index (RefMap | None) – An identity map object -> index to accumulate into. When supplied (e.g. by split_context), indices are global across multiple calls so cross-object references share indices. A fresh map is used if None.

  • treefy_state (bool) – If True the state mapping holds TreefyState values; if False it holds the raw State objects.

Returns:

The static structure and the dynamic state mapping.

Return type:

Tuple[GraphDef, NestedDict]

Raises:

  • TypeError – If ref_index is not a RefMap, or a static attribute value is not hashable.

  • ValueError – If node is not a graph node or pytree (e.g. a bare State).

Examples

>>> import brainstate
>>> model = brainstate.nn.Linear(2, 3)
>>> graphdef, states = brainstate.graph.flatten(model)
>>> isinstance(states, brainstate.util.NestedDict)
True
Contents