Merge branch 'main' into justinchu/ser-initializer

Author: justinchuby

.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"]

README.md

@@ -15,6 +15,16 @@ models using a subset of Python. ONNX Script is:
 * **Debuggable:** allows for eager-mode evaluation that provides for a
   more delightful ONNX model debugging experience.
 
+This repo also covers:
+
+* **ONNX IR:** an in-memory IR that supports the full ONNX spec, designed
+  for graph construction, analysis and transformation.
+* **ONNX Script Optimizer:** provides functionality to optimize an ONNX
+  model by performing optimizations and clean-ups such as constant folding,
+  dead code elimination, etc.
+* **ONNX Rewriter:** provides functionality to replace certain patterns in
+  an ONNX graph with replacement patterns based on user-defined rewrite rules.
+
 Note however that ONNX Script does **not** intend to support the entirety
 of the Python language.
 
@@ -142,6 +152,85 @@ result = Hardmax(v)
 
 More examples can be found in the [docs/examples](docs/examples) directory.
 
+## ONNX IR
+
+An in-memory IR that supports the full ONNX spec, designed for graph construction, analysis and transformation.
+
+### Features
+
+* **Full ONNX spec support:** all valid models representable by ONNX protobuf,
+  and a subset of invalid models (so you can load and fix them).
+* **Low memory footprint:** mmap'ed external tensors; unified interface for
+  ONNX TensorProto, Numpy arrays and PyTorch Tensors etc. No tensor size
+  limitation. Zero copies.
+* **Straightforward access patterns:** Access value information and traverse the
+  graph topology at ease.
+* **Robust mutation:** Create as many iterators as you like on the graph while mutating it.
+* **Speed:** Performant graph manipulation, serialization/deserialization to Protobuf.
+* **Pythonic and familiar APIs:** Classes define Pythonic apis and still map to
+  ONNX protobuf concepts in an intuitive way.
+
+## ONNX Script Tools
+
+### ONNX Optimizer
+
+The ONNX Script Optimizer tool provides the user with the functionality to optimize an ONNX model by performing optimizations and clean-ups such as constant folding, dead code elimination, etc. In order to utilize the optimizer tool:
+
+```python
+import onnxscript
+
+onnxscript.optimizer.optimize(onnx_model)
+```
+
+For a detailed summary of all the optimizations applied by the optimizer call, refer to the tutorial [Optimizing a Model using the Optimizer](https://onnxscript.ai/tutorial/optimizer/optimize.html)
+
+### ONNX Rewriter
+
+The ONNX Rewriter tool provides the user with the functionality to replace certain patterns in an ONNX graph with another pattern based on user-defined rewrite rules. The rewriter tools allows two different methods in which patterns in the graph can be rewritten.
+
+### Pattern-based rewriting
+
+For this style of rewriting, the user provides a `target_pattern` that is to be replaced, a `replacement_pattern` and a `match_condition` (pattern rewrite will occur only if the match condition is satisfied). A simple example on how to use the pattern-based rewriting tool is as follows:
+
+```python
+from onnxscript.rewriter import pattern
+
+# The target pattern
+def erf_gelu_pattern(op, x):
+    return 0.5 * (x * (op.Erf(x / math.sqrt(2)) + 1.0))
+
+def erf_gelu_pattern_2(op, x):
+    return (x * (op.Erf(x / math.sqrt(2)) + 1.0)) * 0.5
+
+# The replacement pattern
+def gelu(op, x: ir.Value):
+    return op.Gelu(x, domain="com.microsoft")
+
+# Create multiple rules
+rule1 = pattern.RewriteRule(
+    erf_gelu_pattern,  # Target Pattern
+    gelu,  # Replacement
+)
+rule2 = pattern.RewriteRule(
+    erf_gelu_pattern_2,  # Target Pattern
+    gelu,  # Replacement
+)
+# Create a Rewrite Rule Set with multiple rules.
+rewrite_rule_set = pattern.RewriteRuleSet([rule1, rule2])
+# Apply rewrites
+model_with_rewrite_applied = onnxscript.rewriter.rewrite(
+    model,  # Original ONNX Model
+    pattern_rewrite_rules=rewrite_rule_set,
+)
+return model_with_rewrite_applied
+```
+
+For a detailed tutorial on how to create target_pattern, replacement_pattern and match_condition blocks in order to utilize the pattern-based rewriter, refer to the tutorial [Pattern-based Rewrite Using Rules](https://onnxscript.ai/tutorial/rewriter/rewrite_patterns.html)
+
+### Function-based rewriting
+
+This style of rewriting matches a `FUNCTION_KEYWORD` and `PACKAGE_NAME` provided by the user to an existing function within the graph and replaces it with a new function provided by the user.
+
 ## Development Guidelines
 
 Every change impacting the converter or the eager evaluation must be

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/ast_utils.py

@@ -8,18 +8,18 @@
 import inspect
 import sys
 import textwrap
-import types
+from typing import Callable
 
 PY_VERSION_GE_39 = sys.version_info >= (3, 9)
 
 
-def get_src_and_ast(f: types.FunctionType) -> tuple[str, ast.FunctionDef]:
+def get_src_and_ast(func: Callable, /) -> tuple[str, ast.FunctionDef]:
     try:
-        src = inspect.getsource(f)
+        src = inspect.getsource(func)
     except OSError as e:
         raise RuntimeError(
             f"Decorator script does not work on dynamically "
-            f"compiled function {f.__name__}."
+            f"compiled function {func.__name__}."
         ) from e
     src = textwrap.dedent(src)
     top_level_ast = ast.parse(src)

onnxscript/_internal/runtime_typing.py

@@ -17,9 +17,11 @@
 T = typing.TypeVar("T", bound=typing.Callable[..., typing.Any])
 
 try:
-    from beartype import beartype as checked
+    from beartype import beartype as _beartype_decorator
     from beartype import roar as _roar
 
+    checked = typing.cast(typing.Callable[[T], T], _beartype_decorator)
+
     # Beartype warns when we import from typing because the types are deprecated
     # in Python 3.9. But there will be a long time until we can move to using
     # the native container types for type annotations (when 3.9 is the lowest

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,
 )