Merge branch 'main' into justinchu/fold-const-shape

Author: justinchuby

.github/workflows/pages.yaml

@@ -42,7 +42,7 @@ jobs:
       - name: Build documentation
         run: python -m sphinx docs dist/html
       - name: Upload documentation archive
-        uses: actions/upload-pages-artifact@v3
+        uses: actions/upload-pages-artifact@v4
         with:
           path: 'dist/html'
       - name: Deploy to GitHub Pages

VERSION

@@ -1 +1 @@
-0.4.0
+0.4.1

noxfile.py

@@ -42,7 +42,7 @@
     "packaging",
     "protobuf",
 )
-ONNX_IR = "onnx_ir==0.1.3"
+ONNX_IR = "onnx_ir==0.1.7"
 ONNX_IR_MAIN = "git+https://github.com/onnx/ir-py.git@main#egg=onnx_ir"
 
 

onnxscript/_framework_apis/torch_2_8.py

@@ -1,6 +1,6 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT License.
-"""Stable APIs for PyTorch 2.7."""
+"""Stable APIs for PyTorch 2.8."""
 
 from __future__ import annotations
 

onnxscript/_framework_apis/torch_2_9.py

@@ -0,0 +1,35 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+"""Stable APIs for PyTorch 2.9."""
+
+from __future__ import annotations
+
+__all__ = [
+    "check_model",
+    "convert_version",
+    "get_torchlib_ops",
+    "optimize",
+    "save_model_with_external_data",
+]
+
+from typing import TYPE_CHECKING
+
+from onnxscript import version_converter
+from onnxscript._framework_apis.torch_2_8 import (
+    check_model,
+    get_torchlib_ops,
+    optimize,
+    save_model_with_external_data,
+)
+
+if TYPE_CHECKING:
+    import onnx_ir as ir
+
+
+def convert_version(model: ir.Model, target_version: int) -> ir.Model:
+    """Convert the model to the specified ONNX opset version.
+
+    Starting from PyTorch 2.9, down conversion is turned on and supported.
+    """
+    version_converter.convert_version(model, target_version, fallback=True)
+    return model

onnxscript/backend/onnx_export.py

@@ -13,6 +13,8 @@
 
 _SINGLE_INDENT = "    "
 
+_SMALL_TENSOR_SIZE = 4
+
 kwlist = {
     "False",
     "None",
@@ -119,7 +121,7 @@ def renamer(name):
 
 def _translate_type(onnx_type):
     """Converts a onnx type into a type defined by *onnxscript*."""
-    return onnxscript.onnx_types.onnx_type_to_onnxscript_repr(onnx_type)
+    return onnxscript.onnx_types.onnx_type_to_onnxscript_repr(onnx_type, reversible=False)
 
 
 def _translate_signature(inputs, outputs):
@@ -350,25 +352,33 @@ def _translate_graph_body(self, graph, opsets, indent=0):
         if hasattr(graph, "initializer"):
             for init in graph.initializer:
                 if self.skip_initializers:
-                    init_py_name = self._translate_onnx_var(init.name)
-                    if init_py_name in self.skipped_initializers:
-                        raise RuntimeError(
-                            f"Initializer {init.name!r} is already present in skipped_initializers."
-                        )
-                    self.skipped_initializers[init_py_name] = init
-                    continue
+                    size = 1
+                    for d in init.dims:
+                        size *= d
+                    if size > _SMALL_TENSOR_SIZE:
+                        init_py_name = self._translate_onnx_var(init.name)
+                        if init_py_name in self.skipped_initializers:
+                            raise RuntimeError(
+                                f"Initializer {init.name!r} is already present in skipped_initializers."
+                            )
+                        self.skipped_initializers[init_py_name] = init
+                        continue
                 node = onnx.helper.make_node(  # noqa: TID251
                     "Constant",
                     [],
                     [self._translate_onnx_var(init.name)],  # type: ignore[list-item]
                     value=init,
                 )
-                code.append(self._translate_node(node, opsets, indent=indent))
+                pyinit = self._translate_node(node, opsets, indent=indent)
+                if pyinit:
+                    code.append(pyinit)
         if hasattr(graph, "sparse_initializer") and len(graph.sparse_initializer) > 0:
             raise NotImplementedError("Unable to convert sparse_initilizer into python.")
         for node in graph.node:
             pynode = self._translate_node(node, opsets, indent=indent)
             if pynode:
+                if node.name:
+                    pynode += f"  # {node.name}"
                 code.append(pynode)
 
         final = "\n".join(code)
@@ -418,7 +428,8 @@ def _translate_attributes(self, node):
     def _translate_if(self, node, opsets, indent=0):
         """Translates a node If into python."""
         sindent = _SINGLE_INDENT * indent
-        code = [f"{sindent}if {node.input[0]}:"]
+        cond = self._translate_onnx_var_ref(node.input[0])
+        code = [f"{sindent}if {cond}:"]
         if len(node.attribute) != 2:
             raise RuntimeError(
                 f"Node {node.op_type!r} expected two attributes not {len(node.attribute)}."
@@ -502,17 +513,21 @@ def _translate_loop(self, node, opsets, indent=0):
 
         rows.extend(self._emit_assign(formal_ins, actual_ins, indent))
 
+        if node.name:
+            node_name = "  # " + node.name
+        else:
+            node_name = ""
         if use_iter_var and not use_loop_cond:
-            rows.append(f"{sindent}for {iter_var} in range({n_iter}):")
+            rows.append(f"{sindent}for {iter_var} in range({n_iter}):{node_name}")
             # The following is a hacky way to suppress the generation of
             # "cond_out = cond_in", which ONNX forces for a FOR loop.
             # TODO: a cleaner solution for this.
             self._name_remappings[-1][cond_out] = self._translate_onnx_var(cond_in)
         elif not use_iter_var and use_loop_cond:
-            rows.append(f"{sindent}while {py_cond}:")
+            rows.append(f"{sindent}while {py_cond}:{node_name}")
         elif use_iter_var and use_loop_cond:
             # TODO: This needs fixing
-            rows.append(f"{sindent}for {iter_var} in range({n_iter}):")
+            rows.append(f"{sindent}for {iter_var} in range({n_iter}):{node_name}")
             rows.append(f"{sindent}{_SINGLE_INDENT}if not {py_cond}:")
             rows.append(f"{sindent}{_SINGLE_INDENT * 2}break")
         else:
@@ -734,11 +749,13 @@ def _substitute_initializers(
 
         def generate_rand(name: str, value: TensorProto) -> str:
             shape = ",".join(str(d) for d in value.dims)
-            if value.data_type != TensorProto.FLOAT:
-                raise NotImplementedError(
-                    f"Unable to generate random initializer for data type {value.data_type}."
-                )
-            return f"{__}{name} = np.random.rand({shape}).astype(np.float32)"
+            if value.data_type == TensorProto.FLOAT:
+                return f"{__}{name} = np.random.rand({shape}).astype(np.float32)"
+            if value.data_type == TensorProto.INT8:
+                return f"{__}{name} = np.random.randint(-128, 127, size=({shape},), dtype=np.int8)"
+            raise NotImplementedError(
+                f"Unable to generate random initializer for data type {value.data_type}."
+            )
 
         random_initializer_values = "\n".join(
             generate_rand(key, value) for key, value in self.skipped_initializers.items()

onnxscript/backend/onnx_export_test.py

@@ -99,6 +99,8 @@ def skip(pattern: str | Pattern, reason: str, *, condition: bool = True):
             "^test_resize_upsample_scales_linear_half_pixel_symmetric",
             "cannot import module, import_module does not work",
         ),
+        # tests are too unstable on Windows, not always the same ones are failing.
+        skip("test_", "cannot import module"),
     )
 
 

onnxscript/onnx_opset/__init__.py

@@ -2,13 +2,11 @@
 # ⚠️ WARNING - AUTO-GENERATED CODE - DO NOT EDIT ⚠️
 # ⚙️ Generated by 'python -m opgen'
 # --------------------------------------------------------------------------
-# Copyright (c) Microsoft Corporation. All rights reserved.
+# Copyright (c) Microsoft Corporation.
 # Licensed under the MIT License.
 # --------------------------------------------------------------------------
 # pylint: disable=W0221,W0222,R0901,W0237
 # mypy: disable-error-code=override
-# ruff: noqa: N801,E741
-# ruff: noqa: D214,D402,D405,D411,D412,D416,D417
 # --------------------------------------------------------------------------
 
 from __future__ import annotations
@@ -40,14 +38,12 @@
 from onnxscript.onnx_opset._impl.opset21 import Opset21
 from onnxscript.onnx_opset._impl.opset22 import Opset22
 from onnxscript.onnx_opset._impl.opset23 import Opset23
+from onnxscript.onnx_opset._impl.opset24 import Opset24
 from onnxscript.onnx_opset._impl.opset_ai_onnx_ml1 import Opset_ai_onnx_ml1
 from onnxscript.onnx_opset._impl.opset_ai_onnx_ml2 import Opset_ai_onnx_ml2
 from onnxscript.onnx_opset._impl.opset_ai_onnx_ml3 import Opset_ai_onnx_ml3
 from onnxscript.onnx_opset._impl.opset_ai_onnx_ml4 import Opset_ai_onnx_ml4
 from onnxscript.onnx_opset._impl.opset_ai_onnx_ml5 import Opset_ai_onnx_ml5
-from onnxscript.onnx_opset._impl.opset_ai_onnx_preview_training1 import (
-    Opset_ai_onnx_preview_training1,
-)
 from onnxscript.values import Opset
 
 __all__ = [
@@ -75,12 +71,12 @@
     "opset21",
     "opset22",
     "opset23",
+    "opset24",
     "opset_ai_onnx_ml1",
     "opset_ai_onnx_ml2",
     "opset_ai_onnx_ml3",
     "opset_ai_onnx_ml4",
     "opset_ai_onnx_ml5",
-    "opset_ai_onnx_preview_training1",
 ]
 
 
@@ -113,12 +109,12 @@
 opset21 = Opset21()
 opset22 = Opset22()
 opset23 = Opset23()
+opset24 = Opset24()
 opset_ai_onnx_ml1 = Opset_ai_onnx_ml1()
 opset_ai_onnx_ml2 = Opset_ai_onnx_ml2()
 opset_ai_onnx_ml3 = Opset_ai_onnx_ml3()
 opset_ai_onnx_ml4 = Opset_ai_onnx_ml4()
 opset_ai_onnx_ml5 = Opset_ai_onnx_ml5()
-opset_ai_onnx_preview_training1 = Opset_ai_onnx_preview_training1()
 all_opsets: Mapping[Tuple[str, int], Opset] = {
     (
         "",
@@ -212,6 +208,10 @@
         "",
         23,
     ): opset23,
+    (
+        "",
+        24,
+    ): opset24,
     (
         "ai.onnx.ml",
         1,
@@ -232,8 +232,4 @@
         "ai.onnx.ml",
         5,
     ): opset_ai_onnx_ml5,
-    (
-        "ai.onnx.preview.training",
-        1,
-    ): opset_ai_onnx_preview_training1,
 }