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

### Changes

- Create env to test with TORCHLIB_EXPERIMENTAL_PREFER_TRACING on
- Test with Python 3.11 as well

[ghstack-poisoned]

Author: justinchuby

onnxscript/function_libs/torch_lib/_flags.py

@@ -39,7 +39,3 @@ 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,7 +99,6 @@ 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.
 
@@ -113,7 +112,6 @@ 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
@@ -130,7 +128,6 @@ 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/tests/function_libs/torch_lib/ops_test_data.py

@@ -501,7 +501,11 @@ def _where_input_wrangler(
     TorchLibOpInfo("acosh", core_ops.aten_acosh),
     TorchLibOpInfo("add", core_ops.aten_add, tolerance={torch.float16: (1e-3, 1e-3)}),
     TorchLibOpInfo("add", core_ops.aten_add_complex, complex=True, trace_only=True),
-    TorchLibOpInfo("addbmm", core_ops.aten_addbmm, tolerance={torch.float32: (2e-5, 2e-5)}),
+    TorchLibOpInfo(
+        "addbmm",
+        core_ops.aten_addbmm,
+        tolerance={torch.float32: (2e-5, 2e-5), torch.float16: (2e-1, 2e-2)},
+    ),
     TorchLibOpInfo("addcdiv", core_ops.aten_addcdiv),
     TorchLibOpInfo("addcmul", core_ops.aten_addcmul, tolerance={torch.float16: (4e-3, 3e-3)}),
     TorchLibOpInfo("addmm", core_ops.aten_addmm)
@@ -522,7 +526,7 @@ def _where_input_wrangler(
         dtypes=(torch.int16, torch.int32, torch.int64),
         reason="ONNX Runtime does not support int inputs to Gemm",
     ),
-    TorchLibOpInfo("addmv", core_ops.aten_addmv),
+    TorchLibOpInfo("addmv", core_ops.aten_addmv, tolerance={torch.float16: (1e-3, 1e-2)}),
     TorchLibOpInfo(
         "addr",
         core_ops.aten_addr,
@@ -640,7 +644,7 @@ def _where_input_wrangler(
             "https://github.com/microsoft/onnxscript/issues/1007"
         ),
     ),
-    TorchLibOpInfo("baddbmm", core_ops.aten_baddbmm),
+    TorchLibOpInfo("baddbmm", core_ops.aten_baddbmm, tolerance={torch.float16: (1e-3, 1e-2)}),
     TorchLibOpInfo("bernoulli", core_ops.aten_bernoulli, nondeterministic=True),
     TorchLibOpInfo(
         # This string is a unique ID. In extra_opinfo.py, we
@@ -845,6 +849,12 @@ def _where_input_wrangler(
         dtypes=(torch.int64, torch.int32),
         reason="fixme: Results do not match with PyTorch. https://github.com/microsoft/onnxscript/issues/854",
     )
+    .xfail(
+        variant_name="tensor_overload",
+        dtypes=(torch.int64, torch.int32, torch.float16),
+        reason="fixme: Results do not match with PyTorch. https://github.com/microsoft/onnxscript/issues/854",
+        enabled_if=not version_utils.torch_older_than("2.1"),
+    )
     .xfail(
         dtypes=(torch.float16,),
         reason="op 'Range' doesn't support float16.",
@@ -884,7 +894,7 @@ def _where_input_wrangler(
         "matmul",
         core_ops.aten_matmul,
         # Windows requires a more relaxed tolerance
-        tolerance={torch.float32: (2e-5, 2e-5)},
+        tolerance={torch.float32: (2e-5, 2e-5), torch.float16: (2e-3, 2e-2)},
     ).skip(
         matcher=lambda sample: torch.numel(sample.input) == 0,
         reason="values of matmul of [m, 0] and [0, n] matrices are undefined",
@@ -1700,7 +1710,12 @@ def _where_input_wrangler(
         variant_name="empty_strides",
         reason="fixme: 'shape' do not match: torch.Size([2, 3, 4, 3]) != torch.Size([2, 3, 4, 2]). https://github.com/microsoft/onnxscript/issues/975",
     ),
-    TorchLibOpInfo("native_batch_norm", core_ops.aten_native_batch_norm, trace_only=True),
+    TorchLibOpInfo(
+        "native_batch_norm",
+        core_ops.aten_native_batch_norm,
+        trace_only=True,
+        tolerance={torch.float16: (9e-3, 7e-4)},
+    ),
     TorchLibOpInfo(
         "ops.aten._native_batch_norm_legit", core_ops.aten_native_batch_norm, trace_only=True
     ),
@@ -1719,9 +1734,11 @@ def _where_input_wrangler(
         "ops.aten.native_group_norm",
         core_ops.aten_native_group_norm,
         trace_only=True,
+        tolerance={torch.float16: (1e-2, 7e-3)},
     ).xfail(
         dtypes=(torch.float16,),
         reason="fixme: 'GroupNormKernelImpl' not implemented for 'Half' in nightly and weekly",
+        enabled_if=version_utils.torch_older_than("2.2"),
     ),
     TorchLibOpInfo(
         "native_layer_norm",
@@ -1809,7 +1826,11 @@ def _where_input_wrangler(
         matcher=lambda sample: len(sample.args) != 1,
         reason="this overload is implemented for bias=None",
     ),
-    TorchLibOpInfo("nn.functional.linear_bias", nn_ops.aten_linear_bias).skip(
+    TorchLibOpInfo(
+        "nn.functional.linear_bias",
+        nn_ops.aten_linear_bias,
+        tolerance={torch.float16: (2e-1, 4e-4)},
+    ).skip(
         # input: input, args: weight, bias; so len(args) == 2 means bias is provided
         matcher=lambda sample: len(sample.args) != 2,
         reason="this overload is implemented for bias!=None",

onnxscript/values.py

@@ -479,9 +479,6 @@ 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",