treefy_states

class brainstate.graph.treefy_states(node, *filters)

Return the treefy state mapping of node, optionally filtered.

Parameters:

• node (Any) – The graph node.

• *filters (type | str | Callable[[Tuple[Key, ...], Any], bool] | bool | EllipsisType | None | Tuple[Filter, ...] | List[Filter]) – Optional filters; when given the mapping is filtered/split.

Returns:

The treefy state mapping (or partitions when multiple filters given).

Return type:

NestedDict | Tuple[NestedDict, ...]