feat(test): use ort 1.14 (#435)

Use ort 1.14 in test

Fix errors in logsumexp and upsample_nearest2d and enable their tests
Parameterized and refactored the backend tests.

Author: justinchuby

.gitignore

@@ -94,6 +94,6 @@ dmypy.json
 onnxscript/tests/models/testoutputs/*
 docs/auto_examples/*
 onnxscript/tests/mylib.onnxlib
-onnxscript/tests/onnx_backend_test_code
+**/onnx_backend_test_code/**
 *.onnxlib
 *.onnx

noxfile.py

@@ -28,7 +28,7 @@
     "pyyaml",
 )
 ONNX = "onnx==1.13"
-ONNX_RUNTIME = "onnxruntime==1.13.1"
+ONNX_RUNTIME = "onnxruntime==1.14"
 PYTORCH = "torch==1.13"
 
 

onnxscript/backend/onnx_backend.py

@@ -6,6 +6,7 @@
 
 import os
 import textwrap
+from typing import Iterator
 
 import numpy
 import onnx
@@ -274,7 +275,7 @@ def to_python(self):
         return final
 
 
-def enumerate_onnx_tests(series, fct_filter=None):
+def enumerate_onnx_tests(series, fct_filter=None) -> Iterator[OnnxBackendTest]:
     """Collects test from a sub folder of `onnx/backend/test`.
     Works as an enumerator to start processing them
     without waiting or storing too much of them.
@@ -285,7 +286,7 @@ def enumerate_onnx_tests(series, fct_filter=None):
         fct_filter: function `lambda testname: boolean` to load or skip
             the test, None for all
 
-    Returns:
+    Yields:
         list of @see cl OnnxBackendTest
     """
     root = os.path.dirname(backend_folder)

onnxscript/backend/onnx_backend_test.py

@@ -0,0 +1,49 @@
+# -------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+# --------------------------------------------------------------------------
+
+import os
+import unittest
+
+import onnxruntime as ort
+
+from onnxscript.backend import onnx_backend
+
+
+def load_function(obj):
+    return ort.InferenceSession(obj.SerializeToString())
+
+
+def run_function(obj, *inputs):
+    names = [i.name for i in obj.get_inputs()]
+    if len(names) < len(inputs):
+        raise AssertionError(f"Got {len(inputs)} inputs but expecting {len(names)}.")
+    feeds = {names[i]: inputs[i] for i in range(len(inputs))}
+    got = obj.run(None, feeds)
+    return got
+
+
+class TestOnnxBackEnd(unittest.TestCase):
+
+    folder = os.path.join(os.path.abspath(os.path.dirname(__file__)), "onnx_backend_test_code")
+
+    def test_enumerate_onnx_tests(self):
+        name = "test_abs"
+        code = list(onnx_backend.enumerate_onnx_tests("node", lambda folder: folder == name))
+        self.assertEqual(len(code), 1)
+
+    def test_enumerate_onnx_tests_run_one(self):
+        done = 0
+        for backend_test in onnx_backend.enumerate_onnx_tests(
+            "node", lambda folder: folder == "test_abs"
+        ):
+            self.assertIn(backend_test.name, repr(backend_test))
+            self.assertGreater(len(backend_test), 0)
+            backend_test.run(load_function, run_function)
+            done += 1
+        self.assertEqual(done, 1)
+
+
+if __name__ == "__main__":
+    unittest.main(verbosity=2)

onnxscript/backend/onnx_export_test.py

@@ -0,0 +1,253 @@
+# -------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+# --------------------------------------------------------------------------
+from __future__ import annotations
+
+import dataclasses
+import importlib
+import pathlib
+import re
+import unittest
+from typing import Pattern
+
+import onnxruntime as ort
+import parameterized
+from onnxruntime.capi import onnxruntime_pybind11_state
+
+import onnxscript
+from onnxscript.backend import onnx_backend, onnx_export
+from onnxscript.tests.models import type_double
+
+
+@dataclasses.dataclass
+class SkipInfo:
+    pattern: Pattern
+    reason: str
+    condition: bool
+
+
+def skip(pattern: str | Pattern, reason: str, *, condition: bool = True):
+    """Create a SkipInfo object.
+
+    Args:
+        pattern: A string or a regular expression to match the ONNX backend test name.
+        reason: The reason why the test is skipped.
+        condition: If False, the test is not skipped.
+    """
+    if isinstance(pattern, str):
+        pattern = re.compile(pattern)
+
+    return SkipInfo(pattern, reason, condition)
+
+
+SKIP_TESTS = (
+    skip(r"_scan_", "Operator Scan is not supported by onnx-script"),
+    skip(r"^test_scan", "Operator Scan is not supported by onnx-script"),
+    skip(
+        r"^test_split",
+        "split has an undefined number of outputs. Current implementation of eager mode is not aware of them",
+    ),
+    skip(
+        r"^test_lstm_defaults",
+        "LSTM has an undefined number of outputs. Current implementation of eager mode is not aware of them",
+    ),
+    skip(
+        r"^test_lstm_with_initial_bias",
+        "LSTM has an undefined number of outputs. Current implementation of eager mode is not aware of them",
+    ),
+    skip(
+        r"^test_lstm_with_peepholes",
+        "LSTM has an undefined number of outputs. Current implementation of eager mode is not aware of them",
+    ),
+    skip(
+        r"^test_optional_get_element_tensor",
+        "ORT Unable to create onnxruntime InferenceSession for executing .OptionalGetElement op with onnx model",
+        condition=ort.__version__ == "1.14.0",
+    ),
+    skip(
+        r"test_loop",
+        "Change when the converter supports support something like 'while i < n and cond:'",
+    ),
+    skip(
+        r"^test_range_float_type_positive_delta_expanded",
+        "Change when the converter supports support something like 'while i < n and cond:'",
+    ),
+    skip(
+        r"^test_range_int32_type_negative_delta_expanded",
+        "Change when the converter supports support something like 'while i < n and cond:'",
+    ),
+)
+
+
+def load_function(obj):
+    return ort.InferenceSession(obj.SerializeToString())
+
+
+def run_function(obj, *inputs):
+    names = [i.name for i in obj.get_inputs()]
+    if len(names) < len(inputs):
+        raise AssertionError(f"Got {len(inputs)} inputs but expecting {len(names)}.")
+    feeds = {names[i]: inputs[i] for i in range(len(inputs))}
+    got = obj.run(None, feeds)
+    return got
+
+
+def extract_functions(name: str, content: str, test_folder: pathlib.Path):
+    """Write the content into a file and import all OnnxFunctions from it."""
+    if not test_folder.exists():
+        test_folder.mkdir()
+        init = test_folder / "__init__.py"
+        init.touch()
+    file = test_folder / f"{name}.py"
+    file.write_text(content, encoding="utf-8")
+
+    import_name = f"onnxscript.tests.{test_folder.parts[-1]}.{name}"
+    try:
+        mod = importlib.import_module(import_name)
+    except (SyntaxError, ImportError) as e:
+        raise AssertionError(
+            f"Unable to import {import_name!r} (file: {file!r})\n----\n{content}"
+        ) from e
+    functions = {
+        k: v for k, v in mod.__dict__.items() if isinstance(v, onnxscript.OnnxFunction)
+    }
+    return functions
+
+
+def exec_main(f, *inputs):
+    output = f(*inputs)
+    if isinstance(output, tuple):
+        return list(output)
+    return [output]
+
+
+class TestOnnxBackEnd(unittest.TestCase):
+
+    test_folder = pathlib.Path(__file__).parent.parent / "tests" / "onnx_backend_test_code"
+
+    def test_export2python(self):
+        proto = type_double.double_abs_subgraph.to_model_proto()
+        code = onnx_export.export2python(proto, rename=True, use_operators=True)
+        self.assertIn("v4 = v2 > v1", code)
+
+    @parameterized.parameterized.expand(  # type: ignore[misc]
+        [
+            (backend_test.name, backend_test)
+            for backend_test in onnx_backend.enumerate_onnx_tests("node")
+        ]
+    )
+    def test_export2python_produces_correct_onnx_script_model(
+        self, _: str, backend_test: onnx_backend.OnnxBackendTest
+    ):
+        for skip_info in SKIP_TESTS:
+            if skip_info.pattern.match(backend_test.name) and skip_info.condition:
+                self.skipTest(skip_info.reason)
+
+        self.assertIn(backend_test.name, repr(backend_test))
+        self.assertGreater(len(backend_test), 0)
+        try:
+            backend_test.run(load_function, run_function)
+        except NotImplementedError as e:
+            self.skipTest(f"Not implemented {e}")
+        except (
+            IndexError,
+            RuntimeError,
+            TypeError,
+            ValueError,
+            AttributeError,
+            onnxruntime_pybind11_state.Fail,  # pylint: disable=c-extension-no-member
+            onnxruntime_pybind11_state.NotImplemented,  # pylint: disable=c-extension-no-member
+            onnxruntime_pybind11_state.InvalidArgument,  # pylint: disable=c-extension-no-member
+        ) as e:
+            self.skipTest(f"Unable to load the model: {e}")
+        except onnxruntime_pybind11_state.RuntimeException as e:  # pylint: disable=c-extension-no-member
+            self.skipTest(f"Unable to run the model: {e}")
+        except AssertionError as e:
+            self.skipTest(f"ORT result mismatches with the expected: {e}")
+
+        code = onnx_export.export2python(
+            backend_test.onnx_model, function_name=f"bck_{backend_test.name}"
+        )
+        self.assertIn("@script()", code)
+        self.assertIn(f"def bck_{backend_test.name}(", code)
+
+        if backend_test.name == "test_resize_downsample_scales_cubic":
+            self.assertIn("Resize(X, None, scales,", code)
+
+        functions = extract_functions(backend_test.name, code, self.test_folder)
+        main_function = functions[f"bck_{backend_test.name}"]
+        self.assertIsNotNone(main_function)
+        proto = main_function.to_model_proto()
+
+        # Opset may be different when an binary operator is used.
+        if backend_test.onnx_model.ir_version != proto.ir_version:
+            if (
+                not backend_test.name.startswith(  # pylint: disable=too-many-boolean-expressions
+                    "test_add"
+                )
+                and not backend_test.name.startswith("test_and")
+                and not backend_test.name.startswith("test_div")
+                and not backend_test.name.startswith("test_equal")
+                and not backend_test.name.startswith("test_greater")
+                and not backend_test.name.startswith("test_less")
+                and not backend_test.name.startswith("test_matmul")
+                and not backend_test.name.startswith("test_mod")
+                and not backend_test.name.startswith("test_mul")
+                and not backend_test.name.startswith("test_not")
+                and not backend_test.name.startswith("test_or")
+                and not backend_test.name.startswith("test_pow")
+                and not backend_test.name.startswith("test_sub")
+                and (backend_test.onnx_model.ir_version, proto.ir_version)
+                not in {(3, 4), (5, 6)}
+            ):
+                # Unexpected behavior for old opsets
+                raise AssertionError(
+                    f"Incompatible ir_version {(backend_test.onnx_model.ir_version)} !="
+                    f" {(proto.ir_version)}\n"
+                    f"{backend_test.onnx_model}\n"
+                    f"-----\n"
+                    f"{proto}"
+                )
+
+        try:
+            session = ort.InferenceSession(proto.SerializeToString())
+        except Exception as e:
+            raise AssertionError(
+                f"Unable to load onnx for test {backend_test.name!r}.\n"
+                f"{onnxscript.proto2text(proto)}\n"
+                f"-----\n"
+                f"{backend_test.onnx_model}"
+            ) from e
+
+        # Check converted onnx
+        def _load_function(_):
+            return session
+
+        def _run_function(obj, *inputs):
+            print("    run ONNX")
+            for i, inp in enumerate(inputs):
+                if inp is None:
+                    print(f"    input {i}: None")
+                else:
+                    print(
+                        f"    input {i}: "
+                        f"dtype={inp.dtype!r} shape={inp.shape!r}"
+                        f"{inp.ravel().tolist()!r}"
+                    )
+            try:
+                return run_function(obj, *inputs)
+            except Exception as e:
+                raise AssertionError(
+                    f"Unable to run test {backend_test.name!r} after conversion.\n"
+                    f"{onnxscript.proto2text(proto)}"
+                ) from e
+
+        backend_test.run(_load_function, _run_function)
+
+        # Check eager mode
+        backend_test.run(lambda _: main_function, exec_main)
+
+
+if __name__ == "__main__":
+    unittest.main(verbosity=2)

onnxscript/function_libs/torch_aten/ops/core.py

@@ -3072,7 +3072,12 @@ def aten_logspace(start: float, end: float, steps: int, base: float = 10.0) -> T
 def aten_logsumexp(self: TReal, dim: INT64, keepdim: int = False) -> TReal:
     """logsumexp(Tensor self, int[1] dim, bool keepdim=False) -> Tensor"""
 
-    return op.ReduceLogSumExp(self, dim, keepdims=keepdim)
+    if op.Size(op.Shape(self)) == 0:
+        # A scalar
+        result = self
+    else:
+        result = op.ReduceLogSumExp(self, dim, keepdims=keepdim)
+    return result
 
 
 def aten_lshift(self: TensorType, other: TensorType) -> TensorType: