walk

nyx.core.paramtree.walk(tree, match, *, prune=None, descend_into_match=False, _prefix=())

Yield (path, node) for the nodes of tree satisfying match.

Depth-first over eqx.Module fields, dict items and list/tuple entries, with the path spelled the way navigate() reads it back.

Parameters:

• tree (pytree) – Root to walk; itself a candidate.

• match (callable) – node -> bool; only matching nodes are yielded.

• prune (callable, optional) – node -> bool; descent stops at a node satisfying it, which is still yielded first if it matches.

• descend_into_match (bool) – Whether to keep walking below a matching node.

Yields:

• tuple – Path from tree to the node.

• object – The node.