unflatten

Contents

unflatten#

class brainstate.graph.unflatten(graphdef, state_mapping, /, *, index_ref=None, index_ref_cache=None)#

Reconstruct a graph node from a GraphDef and a state mapping.

Parameters:

  • graphdef (GraphDef) – The static structure produced by flatten().

  • state_mapping (NestedDict) – The dynamic state mapping (or a merge of filtered mappings).

  • index_ref (dict | None) – A global index -> rebuilt object table to accumulate into. Shared across a merge_context so references resolve across calls.

  • index_ref_cache (dict | None) – A cache of pre-existing objects keyed by global index; when supplied, matching nodes/states are reused and updated in place (the “update” path) instead of created fresh.

Returns:

The reconstructed root object.

Return type:

Any

Raises:

Examples

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