Repository: microsoft/onnxscript

Issue:
Create an option to auto create reproduction scripts for failed ops
For failed ops that needed to be debugged in onnx or ort, it would be nice if we can add an option in the test script to autogen a repro script for each failures.


Discussion:

Resolving patch:
diff --git a/onnxscript/function_libs/torch_lib/ops/core.py b/onnxscript/function_libs/torch_lib/ops/core.py
--- a/onnxscript/function_libs/torch_lib/ops/core.py
+++ b/onnxscript/function_libs/torch_lib/ops/core.py
@@ -735,10 +735,21 @@ def aten_atanh(self: TFloat) -> TFloat:
     return op.Atanh(self)
 
 
-@torch_op("aten::atleast_1d", private=True)
-def _aten_atleast_1d_onnx(self: Sequence[TTensor]) -> TTensor:
+@torch_op("aten::atleast_1d")
+def aten_atleast_1d(self: TTensor) -> TTensor:
     """atleast_1d(Tensor self) -> Tensor"""
 
+    shape = op.Shape(self)
+    rank = op.Size(shape)
+    if rank == 0:
+        self = op.Reshape(self, op.Constant(value_ints=[1]))
+    return self
+
+
+@torch_op("aten::atleast_1d.Sequence")
+def aten_atleast_1d_sequence(self: Sequence[TTensor]) -> TTensor:
+    """atleast_1d.Sequence(Tensor[] tensors) -> Tensor[]"""
+
     @graph()
     def reshape_to_1d(tensor):
         shape = op.Shape(tensor)
@@ -750,25 +761,20 @@ def reshape_to_1d(tensor):
     return op.SequenceMap(self, body=reshape_to_1d)
 
 
-@torch_op("aten::atleast_1d")
-def aten_atleast_1d(self: Sequence[TTensor]) -> TTensor:
-    return _aten_atleast_1d_onnx(self)
-
-
-@torch_op("aten::atleast_1d")
-def aten_atleast_1d_single_tensor(self: TTensor) -> TTensor:
-    """atleast_1d(Tensor self) -> Tensor"""
+@torch_op("aten::atleast_2d")
+def aten_atleast_2d(self: TTensor) -> TTensor:
+    """atleast_2d(Tensor self) -> Tensor"""
 
     shape = op.Shape(self)
     rank = op.Size(shape)
-    if rank == 0:
-        self = op.Reshape(self, op.Constant(value_ints=[1]))
+    if rank <= 1:
+        self = op.Reshape(self, op.Constant(value_ints=[1, -1]))
     return self
 
 
-@torch_op("aten::atleast_2d", private=True)
-def _aten_atleast_2d_onnx(self: Sequence[TTensor]) -> TTensor:
-    """atleast_2d(Tensor self) -> Tensor"""
+@torch_op("aten::atleast_2d.Sequence")
+def aten_atleast_2d_sequence(self: Sequence[TTensor]) -> TTensor:
+    """atleast_2d.Sequence(Tensor[] tensors) -> Tensor[]"""
 
     @graph()
     def reshape_to_2d(tensor):
@@ -781,25 +787,22 @@ def reshape_to_2d(tensor):
     return op.SequenceMap(self, body=reshape_to_2d)
 
 
-@torch_op("aten::atleast_2d")
-def aten_atleast_2d(self: Sequence[TTensor]) -> TTensor:
-    return _aten_atleast_2d_onnx(self)
-
-
-@torch_op("aten::atleast_2d")
-def aten_atleast_2d_single_tensor(self: TTensor) -> TTensor:
-    """atleast_2d(Tensor self) -> Tensor"""
+@torch_op("aten::atleast_3d")
+def aten_atleast_3d(self: TTensor) -> TTensor:
+    """atleast_3d(Tensor self) -> Tensor"""
 
     shape = op.Shape(self)
     rank = op.Size(shape)
     if rank <= 1:
-        self = op.Reshape(self, op.Constant(value_ints=[1, -1]))
+        self = op.Reshape(self, op.Constant(value_ints=[1, -1, 1]))
+    elif rank == 2:
+        self = op.Unsqueeze(self, op.Constant(value_ints=[-1]))
     return self
 
 
-@torch_op("aten::atleast_3d")
-def aten_atleast_3d(self: Sequence[TTensor]) -> TTensor:
-    """atleast_3d(Tensor self) -> Tensor"""
+@torch_op("aten::atleast_3d.Sequence")
+def aten_atleast_3d_sequence(self: Sequence[TTensor]) -> TTensor:
+    """atleast_3d.Sequence(Tensor[] tensors) -> Tensor[]"""
 
     @graph()
     def reshape_to_3d(tensor):
@@ -814,19 +817,6 @@ def reshape_to_3d(tensor):
     return op.SequenceMap(self, body=reshape_to_3d)
 
 
-@torch_op("aten::atleast_3d")
-def aten_atleast_3d_single_tensor(self: TTensor) -> TTensor:
-    """atleast_3d(Tensor self) -> Tensor"""
-
-    shape = op.Shape(self)
-    rank = op.Size(shape)
-    if rank <= 1:
-        self = op.Reshape(self, op.Constant(value_ints=[1, -1, 1]))
-    elif rank == 2:
-        self = op.Unsqueeze(self, op.Constant(value_ints=[-1]))
-    return self
-
-
 @torch_op("aten::baddbmm")
 def aten_baddbmm(
     self: TRealOrUInt8,
@@ -2999,20 +2989,21 @@ def aten_hspmm(mat1: TensorType, mat2: TensorType) -> TensorType:
     raise NotImplementedError()
 
 
-@torch_op("aten::hstack", trace_only=True)
+@torch_op("aten::hstack")
 def aten_hstack(tensors: Sequence[TTensor]) -> TTensor:
     """hstack(Tensor[] tensors) -> Tensor"""
 
-    # Use another onnx function
-    tensors = _aten_atleast_1d_onnx(tensors)
+    @graph()
+    def reshape_to_1d(tensor):
+        shape = op.Shape(tensor)
+        rank = op.Size(shape)
+        if rank == 0:
+            tensor = op.Reshape(tensor, op.Constant(value_ints=[1]))
+        return tensor
+
+    tensors_1d = op.SequenceMap(tensors, body=reshape_to_1d)
 
-    # NOTE: The if/else graph has different shape/type which breaks the
-    #       graph matching. We need to use trace_only.
-    if len(tensors[0].shape) == 1:
-        result = op.ConcatFromSequence(tensors, axis=0, new_axis=0)
-    else:
-        result = op.ConcatFromSequence(tensors, axis=1, new_axis=0)
-    return result
+    return op.ConcatFromSequence(tensors_1d, axis=1, new_axis=0)
 
 
 def aten_hypot(self: TensorType, other: TensorType) -> TensorType:
@@ -7173,8 +7164,18 @@ def aten_view_copy(self: TTensor, size: IntType) -> TTensor:
 def aten_vstack(tensors: Sequence[TTensor]) -> TTensor:
     """vstack(Tensor[] tensors) -> Tensor"""
 
-    tensors = _aten_atleast_2d_onnx(tensors)
-    return op.ConcatFromSequence(tensors, axis=0)
+    # The same logic as atleast_2d duplicated here to keep
+    # the function self contained
+    @graph()
+    def reshape_to_2d(tensor):
+        shape = op.Shape(tensor)
+        rank = op.Size(shape)
+        if rank <= 1:
+            tensor = op.Reshape(tensor, op.Constant(value_ints=[1, -1]))
+        return tensor
+
+    tensors_2d = op.SequenceMap(tensors, body=reshape_to_2d)
+    return op.ConcatFromSequence(tensors_2d, axis=0)
 
 
 @torch_op(("aten::where", "aten::where.self"))


Test patch:
diff --git a/onnxscript/tests/function_libs/torch_lib/error_reproduction.py b/onnxscript/tests/function_libs/torch_lib/error_reproduction.py
new file mode 100644
--- /dev/null
+++ b/onnxscript/tests/function_libs/torch_lib/error_reproduction.py
@@ -0,0 +1,121 @@
+from __future__ import annotations
+
+import pathlib
+import platform
+import sys
+import time
+import traceback
+from typing import Any, Mapping
+
+import numpy as np
+import onnx
+import onnxruntime as ort
+import torch
+
+_REPRODUCTION_TEMPLATE = '''\
+import google.protobuf.text_format
+import numpy as np
+from numpy import array, float16, float32, float64, int32, int64
+import onnx
+import onnxruntime as ort
+
+# Run n times
+N = 1
+
+onnx_model_text = """
+{onnx_model_text}
+"""
+
+ort_inputs = {ort_inputs}
+
+# Set up the inference session
+session_options = ort.SessionOptions()
+session_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_DISABLE_ALL
+onnx_model = onnx.ModelProto()
+google.protobuf.text_format.Parse(onnx_model_text, onnx_model)
+
+# Uncomment this line to save the model to a file for examination
+# onnx.save_model(onnx_model, "{short_test_name}.onnx")
+
+onnx.checker.check_model(onnx_model)
+session = ort.InferenceSession(onnx_model.SerializeToString(), session_options, providers=("CPUExecutionProvider",))
+
+# Run the model
+for _ in range(N):
+    ort_outputs = session.run(None, ort_inputs)
+'''
+
+_ISSUE_MARKDOWN_TEMPLATE = """
+### Summary
+
+ONNX Runtime raises `{error_text}` when executing test `{test_name}` in ONNX Script `TorchLib`.
+
+To recreate this report, use
+
+```bash
+CREATE_REPRODUCTION_REPORT=1 python -m pytest onnxscript/tests/function_libs/torch_lib/ops_test.py -k {short_test_name}
+```
+
+### To reproduce
+
+```python
+{reproduction_code}
+```
+
+### Full error stack
+
+```
+{error_stack}
+```
+
+### Environment
+
+```
+{sys_info}
+```
+"""
+
+
+def create_reproduction_report(
+    test_name: str,
+    onnx_model: onnx.ModelProto,
+    ort_inputs: Mapping[str, Any],
+    error: Exception,
+) -> None:
+    # NOTE: We choose to embed the ONNX model as a string in the report instead of
+    # saving it to a file because it is easier to share the report with others.
+    onnx_model_text = str(onnx_model)
+    with np.printoptions(threshold=sys.maxsize):
+        ort_inputs = dict(ort_inputs.items())
+        input_text = str(ort_inputs)
+    error_text = str(error)
+    error_stack = error_text + "\n" + "".join(traceback.format_tb(error.__traceback__))
+    sys_info = f"""\
+OS: {platform.platform()}
+Python version: {sys.version}
+onnx=={onnx.__version__}
+onnxruntime=={ort.__version__}
+numpy=={np.__version__}
+torch=={torch.__version__}"""
+    short_test_name = test_name.split(".")[-1]
+    reproduction_code = _REPRODUCTION_TEMPLATE.format(
+        onnx_model_text=onnx_model_text,
+        ort_inputs=input_text,
+        short_test_name=short_test_name,
+    )
+
+    markdown = _ISSUE_MARKDOWN_TEMPLATE.format(
+        error_text=error_text,
+        test_name=test_name,
+        short_test_name=short_test_name,
+        reproduction_code=reproduction_code,
+        error_stack=error_stack,
+        sys_info=sys_info,
+    )
+
+    # Turn test name into a valid file name
+    markdown_file_name = f'{short_test_name.replace("/", "-").replace(":", "-")}-{str(time.time()).replace(".", "_")}.md'
+    reports_dir = pathlib.Path("error_reports")
+    reports_dir.mkdir(parents=True, exist_ok=True)
+    with open(reports_dir / markdown_file_name, "w", encoding="utf-8") as f:
+        f.write(markdown)
diff --git a/onnxscript/tests/function_libs/torch_lib/ops_test.py b/onnxscript/tests/function_libs/torch_lib/ops_test.py
--- a/onnxscript/tests/function_libs/torch_lib/ops_test.py
+++ b/onnxscript/tests/function_libs/torch_lib/ops_test.py
@@ -12,6 +12,13 @@
 
     pytest onnxscript/tests/function_libs/torch_lib/ops_test.py -k nn_functional_scaled_dot_product_attention
 
+## Environment variables
+
+1. Set environment variable `CATCH_ORT_SEGFAULT=1` to catch segmentation faults
+in onnxruntime by running the inference sessions in a separate process.
+
+2. Set `CREATE_REPRODUCTION_REPORT=1` to create markdown files for reproduction of
+errors.
 """
 from __future__ import annotations
 
@@ -209,7 +216,8 @@ def run_test_output_match(
                     # TODO(justinchuby): Find a more general solution
                     reference_torch_outputs = [reference_torch_outputs]
 
-                function_output = function_executor(reference_torch_outputs)(
+                test_name = test_suite.id()
+                function_output = function_executor(test_name, reference_torch_outputs)(
                     onnx_function, input_onnx, kwargs_onnx
                 )
                 # Finally we re-flatten everything
diff --git a/onnxscript/tests/function_libs/torch_lib/ops_test_common.py b/onnxscript/tests/function_libs/torch_lib/ops_test_common.py
--- a/onnxscript/tests/function_libs/torch_lib/ops_test_common.py
+++ b/onnxscript/tests/function_libs/torch_lib/ops_test_common.py
@@ -32,6 +32,7 @@
 import onnxscript
 import onnxscript.evaluator
 from onnxscript.function_libs.torch_lib import graph_building
+from onnxscript.tests.function_libs.torch_lib import error_reproduction
 
 T = TypeVar("T")
 
@@ -424,16 +425,17 @@ 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(
-        # 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
-    return TORCH_DTYPE_TO_ONNX_STRING[dtype] in allowed_type_strs
+    # Here we consider seq(tensor(float)) compatible with tensor(float) as well
+    return any(TORCH_DTYPE_TO_ONNX_STRING[dtype] in type_str for type_str in allowed_type_strs)
 
 
 def graph_executor(
+    test_name: str,
     outputs: Sequence[Any],
 ) -> Callable[[Callable[..., Any], tuple[Any], dict[str, Any]], None]:
     """Eagerly executes a function."""
@@ -523,7 +525,10 @@ def _capture_graph_and_evaluate_torch_script_evaluator(function: Callable, args,
             ) from e
 
         try:
-            if os.environ.get("CATCH_ORT_SEGFAULT") == "1":
+            if (
+                os.environ.get("CATCH_ORT_SEGFAULT") == "1"
+                or os.environ.get("CREATE_REPRODUCTION_REPORT") == "1"
+            ):
                 # Use an individual process to run ONNX Runtime to catch segfaults
                 return _safe_ort_session_run(onnx_model.SerializeToString(), ort_inputs)
 
@@ -537,24 +542,41 @@ def _capture_graph_and_evaluate_torch_script_evaluator(function: Callable, args,
             onnxruntime.capi.onnxruntime_pybind11_state.NotImplemented,
             # pylint: enable=c-extension-no-member
         ) as e:
-            raise AssertionError(
+            if os.environ.get("CREATE_REPRODUCTION_REPORT") == "1":
+                error_reproduction.create_reproduction_report(
+                    test_name, onnx_model, ort_inputs, e
+                )
+            raise RuntimeError(
                 "ONNX Runtime failed to evaluate:\n"
                 + _format_model_and_input_information(onnx_model, ort_inputs)
             ) from e
         except OrtAbortedError as e:
-            raise AssertionError(
+            if os.environ.get("CREATE_REPRODUCTION_REPORT") == "1":
+                # Save the model and inputs to a file for reproduction
+                error_reproduction.create_reproduction_report(
+                    test_name, onnx_model, ort_inputs, e
+                )
+            raise OrtAbortedError(
                 "ONNX Runtime aborted:\n"
                 + _format_model_and_input_information(onnx_model, ort_inputs)
             ) from e
+        except Exception as e:
+            if os.environ.get("CREATE_REPRODUCTION_REPORT") == "1":
+                error_reproduction.create_reproduction_report(
+                    test_name, onnx_model, ort_inputs, e
+                )
+            raise
 
     return _capture_graph_and_evaluate_torch_script_evaluator
 
 
 def eager_executor(
+    test_name: str,
     outputs,
 ) -> Callable[[Callable[..., Any], tuple[Any], dict[str, Any]], None]:
     """Eagerly executes a function."""
 
+    del test_name  # Unused
     del outputs  # Unused
 
     def executor(function, args, kwargs):
diff --git a/onnxscript/tests/function_libs/torch_lib/ops_test_data.py b/onnxscript/tests/function_libs/torch_lib/ops_test_data.py
--- a/onnxscript/tests/function_libs/torch_lib/ops_test_data.py
+++ b/onnxscript/tests/function_libs/torch_lib/ops_test_data.py
@@ -519,29 +519,59 @@ def _where_input_wrangler(
     TorchLibOpInfo("atan", core_ops.aten_atan),
     TorchLibOpInfo("atan2", core_ops.aten_atan2),
     TorchLibOpInfo("atanh", core_ops.aten_atanh),
-    TorchLibOpInfo("atleast_1d", core_ops.aten_atleast_1d),
-    TorchLibOpInfo(
-        "atleast_1d_single_tensor",
-        core_ops.aten_atleast_1d_single_tensor,
-    ).skip(
+    TorchLibOpInfo("atleast_1d", core_ops.aten_atleast_1d).skip(
         matcher=lambda sample: isinstance(sample.input, (list, tuple)),
-        reason="atleast_1d_single_tensor overload takes single tensor as input",
+        reason="takes single tensor as input",
     ),
-    TorchLibOpInfo("atleast_2d", core_ops.aten_atleast_2d),
     TorchLibOpInfo(
-        "atleast_2d_single_tensor",
-        core_ops.aten_atleast_2d_single_tensor,
-    ).skip(
+        "atleast_1d_Sequence",
+        core_ops.aten_atleast_1d_sequence,
+    )
+    .skip(
+        matcher=lambda sample: not isinstance(sample.input, (list, tuple)),
+        reason="takes tensor sequences only",
+    )
+    .xfail(
+        reason=(
+            "fixme: [ONNXRuntimeError] : 1 : FAIL : This is an invalid model. Error: Duplicate definition of name (_0x9370ed0_rank)."
+            "https://github.com/microsoft/onnxscript/issues/960"
+        )
+    ),
+    TorchLibOpInfo("atleast_2d", core_ops.aten_atleast_2d).skip(
         matcher=lambda sample: isinstance(sample.input, (list, tuple)),
-        reason="atleast_2d_single_tensor overload takes single tensor as input",
+        reason="takes single tensor as input",
     ),
-    TorchLibOpInfo("atleast_3d", core_ops.aten_atleast_3d),
     TorchLibOpInfo(
-        "atleast_3d_single_tensor",
-        core_ops.aten_atleast_3d_single_tensor,
-    ).skip(
+        "atleast_2d_Sequence",
+        core_ops.aten_atleast_2d_sequence,
+    )
+    .skip(
+        matcher=lambda sample: not isinstance(sample.input, (list, tuple)),
+        reason="takes tensor sequences only",
+    )
+    .xfail(
+        reason=(
+            "fixme: [ONNXRuntimeError] : 1 : FAIL : This is an invalid model. Error: Duplicate definition of name (_0x9370ed0_rank)."
+            "https://github.com/microsoft/onnxscript/issues/960"
+        )
+    ),
+    TorchLibOpInfo("atleast_3d", core_ops.aten_atleast_3d).skip(
         matcher=lambda sample: isinstance(sample.input, (list, tuple)),
-        reason="atleast_3d_single_tensor overload takes single tensor as input",
+        reason="takes single tensor as input",
+    ),
+    TorchLibOpInfo(
+        "atleast_3d_Sequence",
+        core_ops.aten_atleast_3d_sequence,
+    )
+    .skip(
+        matcher=lambda sample: not isinstance(sample.input, (list, tuple)),
+        reason="takes tensor sequences only",
+    )
+    .xfail(
+        reason=(
+            "fixme: [ONNXRuntimeError] : 1 : FAIL : This is an invalid model. Error: Duplicate definition of name (_0x9370ed0_rank)."
+            "https://github.com/microsoft/onnxscript/issues/960"
+        )
     ),
     TorchLibOpInfo("baddbmm", core_ops.aten_baddbmm),
     TorchLibOpInfo("bernoulli", core_ops.aten_bernoulli, nondeterministic=True),
@@ -1270,7 +1300,7 @@ def _where_input_wrangler(
         "vstack",
         core_ops.aten_vstack,
     ).xfail(
-        reason="fixme: A bug of constant-propagation optimization within the subgraph, we can avoid it by turning off graph-optimizations in session options",
+        reason="fixme: [ONNXRuntimeError] : 1 : FAIL : This is an invalid model. Error: Duplicate definition of name (_0x62afb00_rank). https://github.com/microsoft/onnxscript/issues/960",
     ),
     TorchLibOpInfo("where", core_ops.aten_where, input_wrangler=_where_input_wrangler).xfail(
         dtypes=(torch.bool,),
@@ -1410,9 +1440,8 @@ def _where_input_wrangler(
     TorchLibOpInfo(
         "hstack",
         core_ops.aten_hstack,
-        trace_only=True,
     ).xfail(
-        reason="fixme: A bug of constant-propagation optimization within the subgraph, we can avoid it by turning off graph-optimizations in session options",
+        reason="fixme: RUNTIME_EXCEPTION : Exception during initialization: Invalid tensor data type 0. https://github.com/microsoft/onnxscript/issues/960",
     ),
     TorchLibOpInfo(
         "nn.functional.grid_sample",
@@ -1828,9 +1857,9 @@ def _where_input_wrangler(
 ops_test_common.duplicate_opinfo(OPS_DB, "arange", ("arange_start", "arange_start_step"))
 ops_test_common.duplicate_opinfo(OPS_DB, "argmax", ("argmax_dim",))
 ops_test_common.duplicate_opinfo(OPS_DB, "argmin", ("argmin_dim",))
-ops_test_common.duplicate_opinfo(OPS_DB, "atleast_1d", ("atleast_1d_single_tensor",))
-ops_test_common.duplicate_opinfo(OPS_DB, "atleast_2d", ("atleast_2d_single_tensor",))
-ops_test_common.duplicate_opinfo(OPS_DB, "atleast_3d", ("atleast_3d_single_tensor",))
+ops_test_common.duplicate_opinfo(OPS_DB, "atleast_1d", ("atleast_1d_Sequence",))
+ops_test_common.duplicate_opinfo(OPS_DB, "atleast_2d", ("atleast_2d_Sequence",))
+ops_test_common.duplicate_opinfo(OPS_DB, "atleast_3d", ("atleast_3d_Sequence",))
 ops_test_common.duplicate_opinfo(OPS_DB, "cat", ("concat", "concatenate"))
 ops_test_common.duplicate_opinfo(OPS_DB, "clone", ("lift_fresh_copy",))
 ops_test_common.duplicate_opinfo(OPS_DB, "full_like", ("full_like_dtype",))