Merge the various data maps and lists into a unified structure | test(torchlib) (#801)

- Create `TorchLibOpInfo` for putting everything about an op in a single
place.
- Fix `dtype_op_schema_compatible` to account for sequence types.
- Fix `_aten_var_mean_dim_onnx` to support float16
- I removed the allow list in favor of the xfail and skips. In a next PR
we will auto create new xfails and skips when we enable tests for new
data types

Author: justinchuby

noxfile.py

@@ -28,7 +28,7 @@
 )
 ONNX = "onnx==1.14.0"
 ONNX_RUNTIME = "onnxruntime==1.15.0"
-PYTORCH = "torch==2.0.0"
+PYTORCH = "torch==2.0.1"
 ONNX_RUNTIME_NIGHTLY_DEPENDENCIES = (
     "flatbuffers",
     "coloredlogs",

onnxscript/function_libs/torch_lib/ops/core.py

@@ -6669,7 +6669,7 @@ def _aten_var_mean_dim_onnx(
     if correction > 0.0:
         self_shape = op.Shape(self)
         dim_size = op.Gather(self_shape, dim, axis=0)
-        numel_float = op.Cast(op.ReduceProd(dim_size, keepdims=0), to=FLOAT.dtype)
+        numel_float = op.CastLike(op.ReduceProd(dim_size, keepdims=0), self)
         mul = op.Mul(var, numel_float)
         sub = op.Sub(numel_float, correction)
         var = op.Div(mul, sub)

onnxscript/tests/function_libs/torch_lib/ops_test.py

@@ -22,7 +22,6 @@
 import onnx
 import onnxruntime as ort
 import parameterized
-import pytest
 import torch
 from torch.testing._internal import common_device_type
 from torch.testing._internal.opinfo import core as opinfo_core
@@ -33,8 +32,7 @@
 from onnxscript.tests.function_libs.torch_lib import ops_test_common, ops_test_data
 
 # All dtypes will be tested on the generated symbolic functions.
-# complex64 would be flattened to float32.
-# add new dtype in the tuple, and also add the new typpe in OPINFO_FUNCTION_TARGET_DTYPE right after the aten function you are testing
+# complex64 will be flattened to float32.
 TESTED_DTYPES = (
     torch.float16,
     torch.float32,
@@ -90,23 +88,16 @@ def _split_function_and_wrangler(
 
 
 # according to https://pytorch.org/docs/stable/testing.html
-OPINFO_PRECISION_TABLE = {
+OPINFO_PRECISION_TABLE: dict[torch.dtype, tuple[float, float]] = {
     # Tolerance value (rtol, atol)
     # The current most relaxed values are for aten::matmul
     torch.float32: (3.7e-5, 1.8e-4),  # default is 1.3e-6, 1e-5
     torch.float16: (1e-3, 1e-5),  # default is 1e-3, 1e-5
 }
 
 
-def _get_rtol_atol_by_dtype(dtype: torch.dtype) -> tuple(Any, Any):
-    if dtype in OPINFO_PRECISION_TABLE:
-        return OPINFO_PRECISION_TABLE[dtype]
-    return (None, None)
-
-
-def _dtype_is_supported_by_op(op_name: str, dtype: torch.dtype) -> bool:
-    dtype_list = ops_test_data.OPINFO_FUNCTION_TARGET_DTYPE.get(op_name)
-    return dtype in dtype_list
+def _get_rtol_atol_by_dtype(dtype: torch.dtype) -> tuple[Any, Any]:
+    return OPINFO_PRECISION_TABLE.get(dtype, (None, None))
 
 
 class TestFunctionValidity(unittest.TestCase):
@@ -121,8 +112,8 @@ def test_all_script_functions_are_onnx_functions(self):
             if not isinstance(func, onnxscript.OnnxFunction):
                 raise AssertionError(
                     f"'{func}' is not an OnnxFunction. Was it decorated with '@torch_op'? "
-                    "If the function is trace_only, please move it to the "
-                    "'ops_test_data.OPINFO_FUNCTION_MAPPING_TRACE_ONLY' dict."
+                    "If the function is trace_only, please specify trace_only=True "
+                    "in the TorchLibOpInfo entry."
                 )
 
     def test_all_trace_only_functions_are_not_onnx_functions(self):
@@ -131,8 +122,8 @@ def test_all_trace_only_functions_are_not_onnx_functions(self):
             if isinstance(func, onnxscript.OnnxFunction):
                 raise AssertionError(
                     f"'{func.name}' is an OnnxFunction. "
-                    "If the function is not trace_only, please move it to the "
-                    "'ops_test_data.OPINFO_FUNCTION_MAPPING_SCRIPTED' dict."
+                    "If the function is not trace_only, please remove trace_only=True "
+                    "in the TorchLibOpInfo entry."
                 )
 
     @parameterized.parameterized.expand(
@@ -318,9 +309,6 @@ def setUp(self) -> None:
     def test_output_match_opinfo_(
         self, device: str, dtype: torch.dtype, op: opinfo_core.OpInfo
     ):
-        if not _dtype_is_supported_by_op(op.name, dtype):
-            pytest.skip(reason=f"{op.name} cannot support {dtype}")
-
         # Base test method for testing each op with the eager executor, used by instantiate_device_type_tests.
         run_test_output_match(
             self,
@@ -341,7 +329,7 @@ def test_output_match_opinfo_(
         [
             info
             for info in ops_test_data.OPS_DB
-            if info.name in ops_test_data.COMPLEX_TESTED_OPS
+            if info.name in ops_test_data.COMPLEX_FUNCTION_MAPPING
         ],
         allowed_dtypes=COMPLEX_TYPES,
     )
@@ -355,7 +343,7 @@ def test_complex_output_match_opinfo_(
             dtype,
             op,
             ops_test_common.eager_executor,
-            ops_test_data.COMPLEX_FUNCTION_MAPPING_SCRIPTED,
+            ops_test_data.COMPLEX_FUNCTION_MAPPING,
         )
 
 
@@ -383,9 +371,6 @@ def setUp(self) -> None:
     def test_output_match_opinfo_(
         self, device: str, dtype: torch.dtype, op: opinfo_core.OpInfo
     ):
-        if not _dtype_is_supported_by_op(op.name, dtype):
-            pytest.skip(reason=f"{op.name} cannot support {dtype}")
-
         # Base test method for testing each op by running the full ONNX graph.
         run_test_output_match(
             self,
@@ -406,7 +391,7 @@ def test_output_match_opinfo_(
         [
             info
             for info in ops_test_data.OPS_DB
-            if info.name in ops_test_data.COMPLEX_TESTED_OPS
+            if info.name in ops_test_data.COMPLEX_FUNCTION_MAPPING
         ],
         allowed_dtypes=COMPLEX_TYPES,
     )
@@ -420,7 +405,7 @@ def test_complex_output_match_opinfo_(
             dtype,
             op,
             ops_test_common.graph_executor,
-            ops_test_data.COMPLEX_FUNCTION_MAPPING_SCRIPTED,
+            ops_test_data.COMPLEX_FUNCTION_MAPPING,
         )
 
 

onnxscript/tests/function_libs/torch_lib/ops_test_common.py

@@ -424,7 +424,9 @@ def dtype_op_schema_compatible(dtype: torch.dtype, schema: onnx.defs.OpSchema) -
     first_input_type_name = schema.inputs[0].type_str
     # Find the type constraint for the first input by matching the parameter name
     first_input_type_constraint = next(
-        (x for x in schema.type_constraints if x.type_param_str == first_input_type_name), None
+        # Here we consider seq(tensor(float)) compatible with tensor(float) as well
+        (x for x in schema.type_constraints if first_input_type_name in x.type_param_str),
+        None,
     )
     assert first_input_type_constraint is not None
     allowed_type_strs = first_input_type_constraint.allowed_type_strs