Merge branch 'main' into rama/optimizer

.github/workflows/main.yaml

@@ -31,10 +31,8 @@ jobs:
           - py311-onnx-weekly
           - py311-ort-nightly
           - py311-experimental-torchlib-tracing
-          - py311-experimental-torchlib-onnx-ir
           - py310
           - py39
-          - py38
         include:
           - name: py311
             python-version: "3.11"
@@ -45,9 +43,6 @@ jobs:
           - name: py39
             python-version: "3.9"
             nox-tag: test
-          - name: py38
-            python-version: "3.8"
-            nox-tag: test
           - name: py312-torch-nightly
             python-version: "3.12"
             nox-tag: test-torch-nightly
@@ -63,9 +58,6 @@ jobs:
           - name: py311-experimental-torchlib-tracing
             python-version: "3.11"
             nox-tag: test-experimental-torchlib-tracing
-          - name: py311-experimental-torchlib-onnx-ir
-            python-version: "3.11"
-            nox-tag: test-experimental-torchlib-onnx-ir
     runs-on: ${{ matrix.os }}
     steps:
       - uses: actions/checkout@v4
@@ -105,7 +97,7 @@ jobs:
       fail-fast: false
       matrix:
         os: [ubuntu-latest]
-        transformers: ["4.37.2", "4.41.2"]
+        transformers: ["4.37.2", "4.41.2", "4.42.3"]
         torch: ["release", "nightly"]
         python_version: ["3.11"]
         nox-tag: ["test-dort"]

docs/test/test_documentation_examples.py

@@ -34,6 +34,9 @@ def do_test_folder(self, folder):
         if tested == 0:
             raise RuntimeError(f"No example was tested in folder {folder}.")
 
+    @unittest.skipIf(
+        sys.platform != "linux", reason="No need to run the documentation on every OS."
+    )
     def test_documentation_examples(self):
         this = os.path.abspath(os.path.dirname(__file__))
         onxc = os.path.normpath(os.path.join(this, "..", ".."))

noxfile.py

@@ -19,7 +19,7 @@
     'numpy==1.26.4; python_version>="3.9"',
     "packaging",
     "parameterized",
-    "psutil",
+    'psutil; sys_platform != "win32"',
     "pytest-cov",
     "pytest-randomly",
     "pytest-subtests",
@@ -28,13 +28,13 @@
     "pyyaml",
     "types-PyYAML",
     "typing_extensions",
-    "ml_dtypes",
+    "ml-dtypes",
 )
 ONNX = "onnx==1.16"
 ONNX_RUNTIME = "onnxruntime==1.17.1"
 PYTORCH = "torch==2.2.2"
 TORCHVISON = "torchvision==0.17.2"
-TRANSFORMERS = "transformers>=4.37.2"
+TRANSFORMERS = "transformers==4.37.2"
 ONNX_RUNTIME_NIGHTLY_DEPENDENCIES = (
     "flatbuffers",
     "coloredlogs",
@@ -134,27 +134,6 @@ def test_experimental_torchlib_tracing(session):
     )
 
 
-@nox.session(tags=["test-experimental-torchlib-onnx-ir"])
-def test_experimental_torchlib_onnx_ir(session):
-    """Test TorchLib using the ONNX IR to build graphs."""
-    session.install(
-        *COMMON_TEST_DEPENDENCIES,
-        PYTORCH,
-        TORCHVISON,
-        ONNX,
-        *ONNX_RUNTIME_NIGHTLY_DEPENDENCIES,
-    )
-    session.install("-r", "requirements/ci/requirements-ort-nightly.txt")
-    session.install(".", "--no-deps")
-    session.run("pip", "list")
-    session.run(
-        "pytest",
-        "tests/function_libs/torch_lib/ops_test.py",
-        *session.posargs,
-        env={"TORCHLIB_EXPERIMENTAL_USE_IR": "1"},
-    )
-
-
 @nox.session(tags=["test-dort"])
 def test_dort(session):
     """Test the conversion of a couple of models from transformers."""
@@ -163,7 +142,7 @@ def test_dort(session):
     )
     torch_version, transformers_version = session.posargs
 
-    if torch_version == "nighly":
+    if torch_version == "nightly":
         session.install(
             "--pre",
             "torch",

onnxscript/_internal/version_utils.py

@@ -2,6 +2,11 @@
 # Licensed under the MIT License.
 """Version utils for testing."""
 
+from __future__ import annotations
+
+import warnings
+from typing import Callable, Sequence
+
 import packaging.version
 
 
@@ -25,6 +30,19 @@ def torch_older_than(version: str) -> bool:
     )
 
 
+def transformers_older_than(version: str) -> bool | None:
+    """Returns True if the transformers version is older than the given version."""
+    try:
+        import transformers  # pylint: disable=import-outside-toplevel
+    except ImportError:
+        return None
+
+    return (
+        packaging.version.parse(transformers.__version__).release
+        < packaging.version.parse(version).release
+    )
+
+
 def is_onnxruntime_training() -> bool:
     """Returns True if the onnxruntime is onnxruntime-training."""
     try:
@@ -74,3 +92,27 @@ def has_transformers():
         return True  # noqa
     except ImportError:
         return False
+
+
+def ignore_warnings(warns: Warning | Sequence[Warning]) -> Callable:  # type: ignore[arg-type]
+    """Catches warnings.
+
+    Args:
+        warns: warnings to ignore
+
+    Returns:
+        decorated function
+    """
+
+    def wrapper(fct):
+        if warns is None:
+            raise AssertionError(f"warns cannot be None for '{fct}'.")
+
+        def call_f(self):
+            with warnings.catch_warnings():
+                warnings.simplefilter("ignore", warns)  # type: ignore[arg-type]
+                return fct(self)
+
+        return call_f
+
+    return wrapper

onnxscript/backend/onnx_export_test.py

@@ -4,8 +4,10 @@
 
 import dataclasses
 import importlib
+import os
 import pathlib
 import re
+import sys
 import unittest
 from typing import Pattern
 
@@ -89,6 +91,17 @@ def skip(pattern: str | Pattern, reason: str, *, condition: bool = True):
     skip(r"^test_ai_onnx_ml_label_encoder", "ONNX Runtime does not support Opset 21 at 1.17"),
 )
 
+if sys.platform == "win32":
+    SKIP_TESTS = (
+        *SKIP_TESTS,
+        skip(r"^test_gemm_beta", "cannot import module, import_module does not work"),
+        skip(
+            r"^test_averagepool_2d_default",
+            "cannot import module, import_module does not work",
+        ),
+        skip("^test_bitwise_not_3d", "cannot import module, import_module does not work"),
+    )
+
 
 def load_function(obj):
     return ort.InferenceSession(obj.SerializeToString(), providers=("CPUExecutionProvider",))
@@ -106,16 +119,24 @@ def run_function(obj, *inputs):
 def extract_functions(name: str, content: str, test_folder: pathlib.Path):
     if not test_folder.exists():
         test_folder.mkdir(exist_ok=True, parents=True)
-        init = test_folder / "__init__.py"
-        init.touch(exist_ok=True)
-    file = test_folder / f"{name}.py"
-    file.write_text(content, encoding="utf-8")
+        init = str(test_folder / "__init__.py")
+        with open(init, "w", encoding="utf-8") as f:
+            f.write("\n")
+    filename = str(test_folder / f"{name}.py")
+    with open(filename, "w", encoding="utf-8") as f:
+        f.write(content + "\n")
+    assert os.path.exists(
+        filename
+    ), f"{filename!r} ({os.path.abspath(filename)!r} does not exist."
     import_name = f"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}"
+            f"Unable to import {import_name!r} (e={e}) (file: {filename!r}, "
+            f"absolute path: {os.path.abspath(filename)!r}, "
+            f"current folder: {os.getcwd()}"
+            f"\n---- CONTENT --\n{content}"
         ) from e
     functions = {
         k: v for k, v in mod.__dict__.items() if isinstance(v, onnxscript.OnnxFunction)
@@ -265,16 +286,6 @@ 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:

onnxscript/function_libs/torch_lib/_flags.py

@@ -54,4 +54,5 @@ def _load_boolean_flag(
 EXPERIMENTAL_USE_IR: bool = _load_boolean_flag(
     "TORCHLIB_EXPERIMENTAL_USE_IR",
     this_will="use the ONNX IR instead of the PyTorch Graph for graph building",
+    deprecated=True,
 )