Merge branch 'main' of https://github.com/sdpython/onnx-array-api

xadupre · xadupre · commit 4c12efdccd71 · 2024-02-22T14:09:19.000+01:00
diff --git a/CHANGELOGS.rst b/CHANGELOGS.rst
@@ -4,6 +4,7 @@ Change Logs
 0.2.0
 +++++
 
+* :pr:`76`: add a mode to compare models without execution
 * :pr:`75`: add QuickGelu to ExtendedReferenceEvaluator
 * :pr:`71`: adds tools to compare two onnx graphs
 * :pr:`61`: adds function to plot onnx model as graphs
diff --git a/_unittests/ut_reference/test_evaluator_yield.py b/_unittests/ut_reference/test_evaluator_yield.py
@@ -1,6 +1,7 @@
 import unittest
 import numpy as np
 from onnx import TensorProto
+from onnx.checker import check_model
 from onnx.helper import (
     make_function,
     make_graph,
@@ -9,6 +10,7 @@
     make_opsetid,
     make_tensor_value_info,
 )
+from onnx.numpy_helper import from_array
 from onnx.parser import parse_model
 from onnx_array_api.ext_test_case import ExtTestCase
 from onnx_array_api.reference import (
@@ -422,13 +424,13 @@ def test_distance_sequence_str(self):
         text = dc.to_str(s1, s2, align)
         self.assertIn("OUTPUT", text)
         expected = """
-            001=|INPUTfloat322x2ABCDA|INPUTfloat322x2ABCDA
-            002=|INPUTfloat322x2ABCDB|INPUTfloat322x2ABCDB
-            003~|INPUTfloat322x3ABCDX|INPUTfloat322x2ABCDX
-            004-|RESULTfloat322x2CEIOExpH|
-            005=|RESULTfloat322x2CEIOLinearRegrY1|RESULTfloat322x2CEIOLinearRegrY1
-            006~|RESULTfloat322x2CEIOAbsY|RESULTfloat322x3CEIPAbsZ
-            007~|OUTPUTfloat322x2CEIOY|OUTPUTfloat322x2CEIPY
+            001=|INPUTfloat322:2x2ABCDA|INPUTfloat322:2x2ABCDA
+            002=|INPUTfloat322:2x2ABCDB|INPUTfloat322:2x2ABCDB
+            003~|INPUTfloat322:2x3ABCDX|INPUTfloat322:2x2ABCDX
+            004-|RESULTfloat322:2x2CEIOExpH|
+            005=|RESULTfloat322:2x2CEIOLinearRegressioY1|RESULTfloat322:2x2CEIOLinearRegressioY1
+            006~|RESULTfloat322:2x2CEIOAbsY|RESULTfloat322:2x3CEIPAbsZ
+            007~|OUTPUTfloat322:2x2CEIOY|OUTPUTfloat322:2x2CEIPY
         """.replace(
             "            ", ""
         ).strip(
@@ -460,6 +462,68 @@ def test_compare_execution(self):
         self.assertIn("CAAA Constant", text)
         self.assertEqual(len(align), 5)
 
+    def test_no_execution(self):
+        model = make_model(
+            make_graph(
+                [
+                    make_node("Unsqueeze", ["X", "zero"], ["xu1"]),
+                    make_node("Unsqueeze", ["xu1", "un"], ["xu2"]),
+                    make_node("Reshape", ["xu2", "shape1"], ["xm1"]),
+                    make_node("Reshape", ["Y", "shape2"], ["xm2c"]),
+                    make_node("Cast", ["xm2c"], ["xm2"], to=1),
+                    make_node("MatMul", ["xm1", "xm2"], ["xm"]),
+                    make_node("Reshape", ["xm", "shape3"], ["Z"]),
+                ],
+                "dummy",
+                [
+                    make_tensor_value_info("X", TensorProto.FLOAT, [32, 128]),
+                    make_tensor_value_info("Y", TensorProto.FLOAT, [3, 5, 128, 64]),
+                ],
+                [make_tensor_value_info("Z", TensorProto.FLOAT, [3, 5, 32, "N"])],
+                [
+                    from_array(np.array([0], dtype=np.int64), name="zero"),
+                    from_array(np.array([1], dtype=np.int64), name="un"),
+                    from_array(np.array([1, 32, 128], dtype=np.int64), name="shape1"),
+                    from_array(np.array([15, 128, 64], dtype=np.int64), name="shape2"),
+                    from_array(np.array([3, 5, 32, 64], dtype=np.int64), name="shape3"),
+                ],
+            )
+        )
+        check_model(model)
+        res1, res2, align, dc = compare_onnx_execution(model, model, mode="nodes")
+        text = dc.to_str(res1, res2, align)
+        self.assertIn("012 = | NODE", text)
+
+        model2 = make_model(
+            make_graph(
+                [
+                    make_node("Unsqueeze", ["X", "zero"], ["xu1"]),
+                    make_node("Unsqueeze", ["xu1", "un"], ["xu2"]),
+                    make_node("Reshape", ["xu2", "shape1"], ["xm1"]),
+                    make_node("Reshape", ["Y", "shape2"], ["xm2c"]),
+                    make_node("MatMul", ["xm1", "xm2c"], ["xm"]),
+                    make_node("Reshape", ["xm", "shape3"], ["Z"]),
+                ],
+                "dummy",
+                [
+                    make_tensor_value_info("X", TensorProto.FLOAT, [32, 128]),
+                    make_tensor_value_info("Y", TensorProto.FLOAT, [3, 5, 128, 64]),
+                ],
+                [make_tensor_value_info("Z", TensorProto.FLOAT, [3, 5, 32, "N"])],
+                [
+                    from_array(np.array([0], dtype=np.int64), name="zero"),
+                    from_array(np.array([1], dtype=np.int64), name="un"),
+                    from_array(np.array([1, 32, 128], dtype=np.int64), name="shape1"),
+                    from_array(np.array([15, 128, 64], dtype=np.int64), name="shape2"),
+                    from_array(np.array([3, 5, 32, 64], dtype=np.int64), name="shape3"),
+                ],
+            )
+        )
+        check_model(model2)
+        res1, res2, align, dc = compare_onnx_execution(model, model2, mode="nodes")
+        text = dc.to_str(res1, res2, align)
+        self.assertIn("012 = | NODE", text)
+
 
 if __name__ == "__main__":
     unittest.main(verbosity=2)
diff --git a/onnx_array_api/_command_lines_parser.py b/onnx_array_api/_command_lines_parser.py
@@ -20,7 +20,7 @@ def get_main_parser() -> ArgumentParser:
         Selects a command.
         
         'translate' exports an onnx graph into a piece of code replicating it,
-        'compares' compares the execution of two onnx models
+        'compare' compares the execution of two onnx models
         """
         ),
     )
@@ -90,6 +90,13 @@ def get_parser_compare() -> ArgumentParser:
         required=True,
         help="second onnx model",
     )
+    parser.add_argument(
+        "-m",
+        "--mode",
+        choices=["execute", "nodes"],
+        default="execute",
+        help="compare the execution ('execute') or the nodes only ('nodes')",
+    )
     parser.add_argument(
         "-v",
         "--verbose",
@@ -112,8 +119,10 @@ def _cmd_compare(argv: List[Any]):
     args = parser.parse_args(argv[1:])
     onx1 = onnx.load(args.model1)
     onx2 = onnx.load(args.model2)
-    res1, res2, align, dc = compare_onnx_execution(onx1, onx2, verbose=args.verbose)
-    text = dc.to_str(res1, res2, align, column_size=args.column_size)
+    res1, res2, align, dc = compare_onnx_execution(
+        onx1, onx2, verbose=args.verbose, mode=args.mode
+    )
+    text = dc.to_str(res1, res2, align, column_size=int(args.column_size))
     print(text)
 
 
@@ -127,7 +136,7 @@ def main(argv: Optional[List[Any]] = None):
             parser = get_main_parser()
             parser.parse_args(argv)
         else:
-            parsers = dict(translate=get_parser_translate)
+            parsers = dict(translate=get_parser_translate, compare=get_parser_compare)
             cmd = argv[0]
             if cmd not in parsers:
                 raise ValueError(
diff --git a/onnx_array_api/reference/evaluator_yield.py b/onnx_array_api/reference/evaluator_yield.py
