graph_to_tree

Contents

graph_to_tree#

class brainstate.graph.graph_to_tree(may_have_graph_nodes, /, *, prefix=<class 'brainstate.typing.Missing'>, split_fn=<function _default_split_fn>, map_non_graph_nodes=False, check_aliasing=True)[source]#

Convert a pytree that may contain graph nodes into a pure pytree.

Every graph node embedded in the input is replaced (by split_fn, which defaults to wrapping it in a NodeStates) so the result is a plain pytree safe to pass through JAX transformations. Aliasing across the split nodes is collected so it can be restored later by tree_to_graph().

Parameters:

  • may_have_graph_nodes (Any) – A pytree whose leaves may include graph nodes (Node subclasses and registered graph-node types). Passed positionally.

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

  • split_fn (Callable[[SplitContext, tuple[TypeVar(KeyEntry), ...], Any, TypeVar(Leaf)], Any]) – Called as split_fn(ctx, keypath, prefix, leaf) for each graph node (and, when map_non_graph_nodes is true, for every other leaf). The default wraps the node’s split result in a NodeStates.

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

  • check_aliasing (bool) – When True (default), verify that shared/aliased graph nodes are consistent with prefix and raise if not.

Return type:

tuple[PyTree, dict[tuple[TypeVar(KeyEntry), ...], int | Array | ndarray]]

Returns:

  • pytree_out (PyTree) – The input pytree with graph nodes replaced by split_fn outputs.

  • find_states (dict) – A mapping from key-path to the State objects discovered while splitting, used to relink shared state on the way back.

See also

tree_to_graph

The inverse conversion.

Contents