tree_to_graph

class brainstate.graph.tree_to_graph(tree, /, *, prefix=<class 'brainstate.typing.Missing'>, merge_fn=<function _merge_tree_node>, is_node_leaf=<function _is_tree_node>, is_leaf=<function _is_tree_node>, map_non_graph_nodes=False)

Convert a pytree of NodeStates back into graph nodes.

The inverse of graph_to_tree(): each leaf recognised as a packed node (by is_node_leaf) is merged back into a live graph node via merge_fn, restoring the original sharing within a single merge_context().

Parameters:

• tree (Any) – A pytree whose node leaves are NodeStates (or whatever is_node_leaf recognises). Passed positionally.

• prefix (Any) – A prefix pytree broadcast over tree and handed to merge_fn per leaf. Defaults to Missing (no prefix).

• merge_fn (Callable[[MergeContext, tuple[TypeVar(KeyEntry), ...], Any, TypeVar(Leaf)], Any]) – Called as merge_fn(ctx, keypath, prefix, leaf) for each node leaf. The default merges leaf.graphdef with leaf.states.

• is_node_leaf (Callable[[TypeVar(Leaf)], bool]) – Predicate selecting which leaves are merged back into graph nodes. Defaults to “is a NodeStates”.

• is_leaf (Callable[[TypeVar(Leaf)], bool]) – Predicate marking where jax.tree flattening stops. Defaults to “is a NodeStates”.

• map_non_graph_nodes (bool) – When True, also run merge_fn on non-node leaves. Defaults to False.

Returns:

The reconstructed pytree with graph nodes restored in place.

Return type:

Any

See also

graph_to_tree

The inverse conversion.