treefy_split

Contents

treefy_split#

class brainstate.graph.treefy_split(node, *filters)#

Split a graph node into a GraphDef and one or more state mappings.

Parameters:

  • node (Any) – The graph node to split.

  • *filters (type | str | Callable[[Tuple[Key, ...], Any], bool] | bool | EllipsisType | None | Tuple[Filter, ...] | List[Filter]) – Optional predicates/types to partition the state. With no filters the single full state mapping is returned; with filters the state is split accordingly (use ... as the last filter to capture the rest).

Returns:

(graphdef, state) with no filters, or (graphdef, state1, state2, ...) with filters.

Return type:

Tuple[GraphDef, Any]

Examples

>>> import brainstate
>>> model = brainstate.nn.Linear(2, 3)
>>> graphdef, params = brainstate.graph.treefy_split(model)
>>> rebuilt = brainstate.graph.treefy_merge(graphdef, params)
>>> isinstance(rebuilt, brainstate.nn.Linear)
True
Contents