Skip incompatible dtypes in op_test using OpSchema | test(torchlib) (#700)

This requires ONNX 1.14. I am skipping the op test for ONNX<1.14
altogether.

I tested locally to make sure no float32 test is skipped with this
mechanism, and that it skips properly according to the type annotation.

Currently no new dtypes are enabled because there are too many failures.

Requires 
- #701 
- #698

Author: justinchuby

onnxscript/tests/function_libs/torch_lib/ops_test.py

@@ -162,6 +162,11 @@ def run_test_output_match(
     # An example is nn.functional.upsample_nearest2d, which has a different signature
     # than the aten operator upsample_nearest2d
     onnx_function, input_wrangler = _split_function_and_wrangler(onnx_function_and_wrangler)
+    if not ops_test_common.dtype_op_schema_compatible(dtype, onnx_function.op_schema):
+        test_suite.skipTest(
+            f"dtype '{dtype}' is not supported by the op '{op.name}'. "
+            f"Type constraints: {onnx_function.op_schema.type_constraints}"
+        )
 
     for i, cpu_sample in enumerate(samples):
         inputs = (cpu_sample.input, *cpu_sample.args)
@@ -251,6 +256,10 @@ def run_test_output_match(
                         raise
 
 
+@unittest.skipIf(
+    version_utils.onnx_older_than("1.14"),
+    "OpSchema not available for functions before ONNX 1.14",
+)
 class TestOutputConsistencyEager(unittest.TestCase):
     """Test output consistency between the ONNX op run with ONNX eager mode and PyTorch eager mode.
 
@@ -279,6 +288,10 @@ def test_output_match_opinfo_(
         run_test_output_match(self, device, dtype, op, ops_test_common.eager_executor)
 
 
+@unittest.skipIf(
+    version_utils.onnx_older_than("1.14"),
+    "OpSchema not available for functions before ONNX 1.14",
+)
 class TestOutputConsistencyFullGraph(unittest.TestCase):
     """Test output consistency between exported ONNX op run as a graph and PyTorch eager mode.
 

onnxscript/tests/function_libs/torch_lib/ops_test_common.py

@@ -355,6 +355,79 @@ def _format_model_and_input_information(onnx_model, inputs):
     )
 
 
+TORCH_DTYPE_TO_ONNX_STRING = {
+    torch.bool: "tensor(bool)",
+    torch.uint8: "tensor(uint8)",
+    torch.int8: "tensor(int8)",
+    torch.int16: "tensor(int16)",
+    torch.int32: "tensor(int32)",
+    torch.int64: "tensor(int64)",
+    torch.float16: "tensor(float16)",
+    torch.float32: "tensor(float)",
+    torch.float64: "tensor(double)",
+    torch.complex64: "tensor(complex64)",
+    torch.complex128: "tensor(complex128)",
+    torch.bfloat16: "tensor(bfloat16)",
+}
+
+
+def dtype_op_schema_compatible(dtype: torch.dtype, schema: onnx.defs.OpSchema) -> bool:
+    """Checks if the dtype is compatible with the schema.
+
+    When a dtype is "compatible" with the schema, it means we can use the dtype
+    to create sample inputs by OpInfo to test the ONNX function and expect outputs to match.
+
+    Args:
+        dtype: The torch dtype used to create sample inputs by OpInfo.
+        schema: The ONNX schema of the function.
+
+    Returns:
+        True if the dtype is compatible with the schema.
+    """
+    if not schema.inputs:
+        # If there are no inputs, we can't check compatibility. Assume it is compatible.
+        # e.g. aten_randn has only attributes.
+        return True
+    if schema.inputs[0].name not in {"self", "input"}:
+        # If the name of the first input is not "self" or "input",
+        # it is usually an input that is not of the same type as the output.
+        # We assume support in this case.
+        #
+        # For example, `aten_ones(size: IntType, dtype: int = FLOAT.dtype)`
+        # has the first input as `size`, which is an integer, but it can support
+        # any dtype.
+        return True
+
+    # Otherwise we check the type constraints of the first input.
+    # For example, when dtype=torch.float32, and the op being tested has the schema
+    # ```
+    # OpSchema(
+    #     name='aten_abs',
+    #     domain='onnxscript.atenlib',
+    #     since_version=1,
+    #     doc='abs(Tensor self) -> Tensor',
+    #     type_constraints=[OpSchema.TypeConstraintParam(type_param_str='TReal', allowed_type_strs=['tensor(float)', 'tensor(int8)', 'tensor(int16)', 'tensor(int32)', 'tensor(int64)', 'tensor(float16)', 'tensor(double)', 'tensor(bfloat16)'], description='')],
+    #     inputs=[OpSchema.FormalParameter(name='self', type_str='TReal', description='', param_option=<FormalParameterOption.Single: 0>, is_homogeneous=True, min_arity=1, differentiation_category=<DifferentiationCategory.Unknown: 0>)],
+    #     outputs=[OpSchema.FormalParameter(name='return_val', type_str='TReal', description='', param_option=<FormalParameterOption.Single: 0>, is_homogeneous=True, min_arity=1, differentiation_category=<DifferentiationCategory.Unknown: 0>)],
+    #     attributes={}
+    # )
+    # ```
+    # we see the first input type is "TReal", corresponding to the type constraint
+    # with allowed types ['tensor(float)', 'tensor(int8)', 'tensor(int16)',
+    # 'tensor(int32)', 'tensor(int64)', 'tensor(float16)', 'tensor(double)',
+    # 'tensor(bfloat16)'].
+    # Since torch.float32 (tensor(float)) is in the allowed types, we return True.
+
+    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
+    )
+    assert first_input_type_constraint is not None
+    allowed_type_strs = first_input_type_constraint.allowed_type_strs
+    return TORCH_DTYPE_TO_ONNX_STRING[dtype] in allowed_type_strs
+
+
 def graph_executor(
     outputs: Sequence[Any],
 ) -> Callable[[Callable[..., Any], tuple[Any], dict[str, Any]], None]:

requirements/ci/requirements-pytorch-nightly.txt

@@ -1,3 +1,3 @@
 --index-url=https://download.pytorch.org/whl/nightly/cpu
 --pre
-torch==2.1.0.dev20230418
+torch