add method check_order

xadupre · xadupre · commit 6532733c6465 · 2023-12-21T14:35:07.000+01:00
diff --git a/_unittests/ut_graph_api/test_graph_builder_optim.py b/_unittests/ut_graph_api/test_graph_builder_optim.py
@@ -1,6 +1,7 @@
 import os
 import unittest
 import onnx
+from onnx.inliner import inline_local_functions
 from onnx_array_api.ext_test_case import ExtTestCase
 from onnx_array_api.graph_api.graph_builder import GraphBuilder
 
@@ -54,7 +55,7 @@ def test_keep_unused_outputs(self):
         self.assertEqual(len(onx.graph.node), 2)
         self.assertEqual(onx.graph.node[0].op_type, "Split")
 
-    def test_check_files(self):
+    def test_check_afiles(self):
         import onnxruntime
 
         data = os.path.join(os.path.dirname(__file__), "data")
@@ -66,8 +67,14 @@ def test_check_files(self):
                     os.path.join(data, f), providers=["CPUExecutionProvider"]
                 )
                 assert sess
-                g = GraphBuilder(onx)
-                g.optimize()
+                onxi = inline_local_functions(onx)
+                sess = onnxruntime.InferenceSession(
+                    onxi.SerializeToString(), providers=["CPUExecutionProvider"]
+                )
+                assert sess
+                g = GraphBuilder(onxi)
+                g.optimize(check_order=True)
+                g.check_order()
                 onx2 = g.to_onnx()
                 sess2 = onnxruntime.InferenceSession(
                     onx2.SerializeToString(), providers=["CPUExecutionProvider"]
diff --git a/onnx_array_api/graph_api/graph_builder.py b/onnx_array_api/graph_api/graph_builder.py
@@ -1,5 +1,5 @@
 from functools import partial
-from typing import Any, Dict, List, Optional, Sequence, Tuple, Union
+from typing import Any, Dict, List, Optional, Sequence, Set, Tuple, Union
 import numpy as np
 import onnx.helper as oh
 import onnx.numpy_helper as onh
@@ -604,14 +604,56 @@ def to_onnx(
         model = oh.make_model(graph, opset_imports=opsets)
         return model
 
-    def optimize(self):
+    def _check_order_node(self, ind: int, node: NodeProto, existing: Set[str]):
+        for i in node.input:
+            if i not in existing:
+                raise RuntimeError(
+                    f"Unknown input {i!r} from node {ind}:{node.op_type}:{node.name}. "
+                    f"Known: {existing}."
+                )
+        for att in node.attribute:
+            if att.type == AttributeProto.GRAPH and att.g:
+                g_existing = existing.copy()
+                for i in att.g.input:
+                    g_existing.add(i.name)
+                for ind2, node2 in enumerate(att.g.node):
+                    self._check_order_node((ind, ind2), node2, g_existing)
+                for o in att.g.output:
+                    if o.name not in g_existing:
+                        raise RuntimeError(
+                            f"Unknown output {o.name!r}. Known: {g_existing}."
+                        )
+        for o in node.output:
+            existing.add(o)
+
+    def check_order(self):
+        existing = set(self.initializers_dict)
+        for i in self.inputs:
+            existing.add(i.name)
+        for ind, node in enumerate(self.nodes):
+            self._check_order_node(ind, node, existing)
+        for o in self.outputs:
+            if o.name not in existing:
+                raise RuntimeError(f"Unknown output {o.name!r}. Known: {existing}.")
+
+    def optimize(self, check_order: bool = False):
+        if check_order:
+            self.check_order()
         self.remove_identity_nodes()
+        if check_order:
+            self.check_order()
         if self.optimization_options.remove_unused:
             self.remove_unused()
+            if check_order:
+                self.check_order()
         if self.optimization_options.constant_folding:
             self.constant_folding()
+            if check_order:
+                self.check_order()
             if self.optimization_options.remove_unused:
                 self.remove_unused()
+                if check_order:
+                    self.check_order()
 
     def remove_unused(self):
         """
