iter_leaf

Contents

iter_leaf#

class brainstate.graph.iter_leaf(node, allowed_hierarchy=(0, 2147483647))#

Iterate (path, value) over every leaf in the graph.

Repeated containers are visited only once (by identity). State leaves are likewise deduplicated by identity: if the same State object is reachable via multiple paths, only its first pre-order occurrence is yielded.

Parameters:

  • node (Any) – The root graph node or pytree.

  • allowed_hierarchy (tuple[int, int]) – (lo, hi) graph-node depth bounds; leaves outside the range are skipped. Depth increments only when descending into a graph node.

Yields:

tuple of (PathParts, Any) – The path to, and value of, each leaf.

Examples

>>> import brainstate
>>> model = brainstate.nn.Linear(2, 3)
>>> leaves = list(brainstate.graph.iter_leaf(model))
>>> all(len(p) >= 1 for p, _ in leaves)
True
Contents