Create env to test with TORCHLIB_EXPERIMENTAL_PREFER_TRACING on | test(torchlib)

.github/workflows/main.yaml

@@ -31,6 +31,8 @@ jobs:
           - py310-torch-nightly
           - py310-onnx-weekly
           - py310-ort-nightly
+          - py311-ort-nightly
+          - py310-experimental-torchlib-tracing
         include:
           - name: py310
             python-version: "3.10"
@@ -50,6 +52,12 @@ jobs:
           - name: py310-ort-nightly
             python-version: "3.10"
             nox-tag: test-ort-nightly
+          - name: py311-ort-nightly
+            python-version: "3.11"
+            nox-tag: test-ort-nightly
+          - name: py310-experimental-torchlib-tracing
+            python-version: "3.10"
+            nox-tag: test-experimental-torchlib-tracing
     runs-on: ${{ matrix.os }}
     steps:
       - uses: actions/checkout@v4

noxfile.py

@@ -11,7 +11,7 @@
 
 COMMON_TEST_DEPENDENCIES = (
     "jinja2",
-    "numpy==1.23.5",
+    "numpy==1.24.4",
     "typing_extensions",
     "beartype!=0.16.0",
     "types-PyYAML",
@@ -95,3 +95,20 @@ def test_ort_nightly(session):
     session.install(".", "--no-deps")
     session.run("pip", "list")
     session.run("pytest", "onnxscript", *session.posargs)
+
+
+@nox.session(tags=["test-experimental-torchlib-tracing"])
+def test_experimental_torchlib_tracing(session):
+    """Test TorchLib with the experimental TORCHLIB_EXPERIMENTAL_PREFER_TRACING flag on."""
+    session.install(
+        *COMMON_TEST_DEPENDENCIES, PYTORCH, 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",
+        "onnxscript/tests/function_libs/torch_lib/ops_test.py",
+        *session.posargs,
+        env={"TORCHLIB_EXPERIMENTAL_PREFER_TRACING": "1"},
+    )

onnxscript/function_libs/torch_lib/_flags.py

@@ -39,3 +39,7 @@ def _load_boolean_flag(
     "TORCHLIB_EXPERIMENTAL_INITIALIZERS_AS_INPUTS",
     this_will="make initializers as inputs to the model graph",
 )
+EXPERIMENTAL_PREFER_TRACING: bool = _load_boolean_flag(
+    "TORCHLIB_EXPERIMENTAL_PREFER_TRACING",
+    this_will="trace all traceable functions to fold if branches and collapse constant expressions",
+)

onnxscript/function_libs/torch_lib/registration.py

@@ -99,6 +99,7 @@ def torch_op(
     trace_only: bool = False,
     private: bool = False,
     complex: bool = False,
+    traceable: bool = False,
 ) -> Callable[[FunctionType], onnxscript.OnnxFunction | onnxscript.values.TracedOnnxFunction]:
     """Register a torch op.
 
@@ -112,6 +113,7 @@ def torch_op(
         private: Whether the function is private (not directly exposed). It should
             be true for all functions with names starting with "_".
         complex: Whether the function expects complex-valued inputs.
+        traceable: Whether the function can be traced.
     """
     if registry is None:
         registry = default_registry
@@ -128,6 +130,7 @@ def wrapper(
         else:
             assert isinstance(func, FunctionType)
             processed_func = onnxscript.script(opset=custom_opset)(func)
+            processed_func.experimental_traceable = traceable
 
         assert registry is not None
         for name_ in _check_and_normalize_names(name):

onnxscript/values.py

@@ -479,6 +479,9 @@ def __init__(
         self._param_schemas: Optional[tuple[ParamSchema, ...]] = None
         self._op_schema: Optional[onnx.defs.OpSchema] = None
 
+        # Experimental fields
+        self.experimental_traceable = False
+
     @property
     @deprecation.deprecated(
         since="0.1",