# Copyright (c) ONNX Project Contributors# SPDX-License-Identifier: Apache-2.0"""Utilities for traversing the IR graph."""from__future__importannotations__all__=["RecursiveGraphIterator",]fromcollections.abcimportIterator,ReversiblefromtypingimportCallable,Unionfromtyping_extensionsimportSelffromonnx_irimport_core,_enumsGraphLike=Union[_core.Graph,_core.Function,_core.GraphView]
[docs]def__init__(self,graph_like:GraphLike,*,recursive:Callable[[_core.Node],bool]|None=None,reverse:bool=False,enter_graph:Callable[[GraphLike],None]|None=None,exit_graph:Callable[[GraphLike],None]|None=None,):"""Iterate over the nodes in the graph, recursively visiting subgraphs. This iterator allows for traversing the nodes of a graph and its subgraphs in a depth-first manner. It supports optional callbacks for entering and exiting subgraphs, as well as a callback `recursive` to determine whether to visit subgraphs contained within nodes. .. versionadded:: 0.1.6 Added the `enter_graph` and `exit_graph` callbacks. Args: graph_like: The graph to traverse. recursive: A callback that determines whether to recursively visit the subgraphs contained in a node. If not provided, all nodes in subgraphs are visited. reverse: Whether to iterate in reverse order. enter_graph: An optional callback that is called when entering a subgraph. exit_graph: An optional callback that is called when exiting a subgraph. """self._graph=graph_likeself._recursive=recursiveself._reverse=reverseself._iterator=self._recursive_node_iter(graph_like)self._enter_graph=enter_graphself._exit_graph=exit_graph