Implement all_nodes() and subgraphs() as convenience methods (#76)

Suggested by @xiaoyu-work, this PR created convenience methods
all_nodes() and subgraphs() for easier access to the nested nodes and
subgraphs. They are created so that it is easier to discover the
methods.

---------

Signed-off-by: Justin Chu <[email protected]>

Author: justinchuby

src/onnx_ir/_core.py

@@ -2366,6 +2366,28 @@ def num_nodes(self) -> int:
         # NOTE: This is a method specific to Graph, not required by the protocol unless proven
         return len(self)
 
+    def all_nodes(self) -> Iterator[Node]:
+        """Get all nodes in the graph and its subgraphs in O(#nodes + #attributes) time.
+
+        This is an alias for ``onnx_ir.traversal.RecursiveGraphIterator(graph)``.
+        Consider using
+        :class:`onnx_ir.traversal.RecursiveGraphIterator` for more advanced
+        traversals on nodes.
+        """
+        # NOTE: This is a method specific to Graph, not required by the protocol unless proven
+        return onnx_ir.traversal.RecursiveGraphIterator(self)
+
+    def subgraphs(self) -> Iterator[Graph]:
+        """Get all subgraphs in the graph in O(#nodes + #attributes) time."""
+        seen_graphs: set[Graph] = set()
+        for node in onnx_ir.traversal.RecursiveGraphIterator(self):
+            graph = node.graph
+            if graph is self:
+                continue
+            if graph is not None and graph not in seen_graphs:
+                seen_graphs.add(graph)
+                yield graph
+
     # Mutation methods
     def append(self, node: Node, /) -> None:
         """Append a node to the graph in O(1) time.
@@ -2871,7 +2893,7 @@ def graphs(self) -> Iterable[Graph]:
         """Get all graphs and subgraphs in the model.
 
         This is a convenience method to traverse the model. Consider using
-        `onnx_ir.traversal.RecursiveGraphIterator` for more advanced
+        :class:`onnx_ir.traversal.RecursiveGraphIterator` for more advanced
         traversals on nodes.
         """
         # NOTE(justinchuby): Given
@@ -2880,11 +2902,8 @@ def graphs(self) -> Iterable[Graph]:
         # (3) Users familiar with onnxruntime optimization tools expect this method
         # I created this method as a core method instead of an iterator in
         # `traversal.py`.
-        seen_graphs: set[Graph] = set()
-        for node in onnx_ir.traversal.RecursiveGraphIterator(self.graph):
-            if node.graph is not None and node.graph not in seen_graphs:
-                seen_graphs.add(node.graph)
-                yield node.graph
+        yield self.graph
+        yield from self.graph.subgraphs()
 
 
 class Function(_protocols.FunctionProtocol, Sequence[Node], _display.PrettyPrintable):
@@ -3018,6 +3037,28 @@ def meta(self) -> _metadata.MetadataStore:
     def metadata_props(self) -> dict[str, str]:
         return self._graph.metadata_props
 
+    def all_nodes(self) -> Iterator[Node]:
+        """Get all nodes in the graph and its subgraphs in O(#nodes + #attributes) time.
+
+        This is an alias for ``onnx_ir.traversal.RecursiveGraphIterator(graph)``.
+        Consider using
+        :class:`onnx_ir.traversal.RecursiveGraphIterator` for more advanced
+        traversals on nodes.
+        """
+        # NOTE: This is a method specific to Graph, not required by the protocol unless proven
+        return onnx_ir.traversal.RecursiveGraphIterator(self)
+
+    def subgraphs(self) -> Iterator[Graph]:
+        """Get all subgraphs in the function in O(#nodes + #attributes) time."""
+        seen_graphs: set[Graph] = set()
+        for node in onnx_ir.traversal.RecursiveGraphIterator(self):
+            graph = node.graph
+            if graph is self._graph:
+                continue
+            if graph is not None and graph not in seen_graphs:
+                seen_graphs.add(graph)
+                yield graph
+
     # Mutation methods
     def append(self, node: Node, /) -> None:
         """Append a node to the function in O(1) time."""

src/onnx_ir/_core_test.py

@@ -1276,6 +1276,72 @@ def test_topological_sort_subgraph(self):
             ("d", "c", "b", "a", ">", "if"),
         )
 
+    def test_all_nodes_returns_all_nodes(self):
+        # Create a graph with a subgraph
+        v0 = _core.Value(name="v0")
+        v1 = _core.Value(name="v1")
+        node0 = _core.Node("", "A", inputs=(v0,), num_outputs=1)
+        node1 = _core.Node("", "B", inputs=(v1,), num_outputs=1)
+        sub_node = _core.Node(
+            "", "Sub", inputs=(node0.outputs[0], node1.outputs[0]), num_outputs=1
+        )
+        subgraph = _core.Graph(
+            inputs=(), outputs=(sub_node.outputs[0],), nodes=(sub_node,), name="subgraph"
+        )
+        main_node = _core.Node(
+            "",
+            "If",
+            inputs=(node0.outputs[0],),
+            attributes=[ir.AttrGraph("then_branch", subgraph)],
+        )
+        graph = _core.Graph(
+            inputs=(v0, v1),
+            outputs=(main_node.outputs[0],),
+            nodes=(node0, node1, main_node),
+            name="main_graph",
+        )
+        all_nodes = list(graph.all_nodes())
+        # Should include node0, node1, main_node, and sub_node
+        self.assertIn(node0, all_nodes)
+        self.assertIn(node1, all_nodes)
+        self.assertIn(main_node, all_nodes)
+        self.assertIn(sub_node, all_nodes)
+        self.assertEqual(len(all_nodes), 4)
+
+    def test_subgraphs_returns_all_subgraphs(self):
+        # Create a graph with two subgraphs
+        v0 = _core.Value(name="v0")
+        v1 = _core.Value(name="v1")
+        node0 = _core.Node("", "A", inputs=(v0,), num_outputs=1)
+        node1 = _core.Node("", "B", inputs=(v1,), num_outputs=1)
+        sub_node1 = _core.Node("", "Sub1", inputs=(node0.outputs[0],), num_outputs=1)
+        sub_node2 = _core.Node("", "Sub2", inputs=(node1.outputs[0],), num_outputs=1)
+        subgraph1 = _core.Graph(
+            inputs=(), outputs=(sub_node1.outputs[0],), nodes=(sub_node1,), name="subgraph1"
+        )
+        subgraph2 = _core.Graph(
+            inputs=(), outputs=(sub_node2.outputs[0],), nodes=(sub_node2,), name="subgraph2"
+        )
+        main_node = _core.Node(
+            "",
+            "If",
+            inputs=(node0.outputs[0],),
+            attributes=[
+                ir.AttrGraph("then_branch", subgraph1),
+                ir.AttrGraph("else_branch", subgraph2),
+            ],
+        )
+        graph = _core.Graph(
+            inputs=(v0, v1),
+            outputs=(main_node.outputs[0],),
+            nodes=(node0, node1, main_node),
+            name="main_graph",
+        )
+        subgraphs = list(graph.subgraphs())
+        self.assertIn(subgraph1, subgraphs)
+        self.assertIn(subgraph2, subgraphs)
+        self.assertEqual(len(subgraphs), 2)
+
 
 class GraphContainersTest(unittest.TestCase):
     """Test containers for input, output and initializers of a graph."""