Update base for Update on "Implement the experimental evaluator for folding branches and castlikes | feat(torchlib)"

As an effort described in #1095, this PR

- Implements the experimental evaluator for folding branches and castlikes so that they are eagerly evaluated when possible.
- Updates implementation for `addr` for it to be traceable.
- Conditionally enabled previously xfailed tests.

Set `TORCHLIB_EXPERIMENTAL_PREFER_TRACING=1` and tested in CI.

E.g. clamp_min now becomes

```
<
   ir_version: 8,
   opset_import: ["" : 18, "pkg.onnxscript.torch_lib.common" : 1],
   producer_name: "pytorch",
   producer_version: "2.2.0"
>
main_graph (int32[5,10,5] input_0, int32[10,5] input_1) => (int32[5,10,5] _val_9) 
   <int64 _val_2, int64[2] _val_3, int64 _val_4, int64 _val_5, bool _val_6, int64 _val_7, bool _val_8>
{
   _val_2 = Size (input_0)
   _val_3 = Shape <start: int = 0> (input_1)
   _val_4 = Size (_val_3)
   _val_5 = Constant <value: tensor = int64 {0}> ()
   _val_6 = Equal (_val_2, _val_5)
   _val_7 = Constant <value: tensor = int64 {0}> ()
   _val_8 = Equal (_val_4, _val_7)
   _val_9 = Max (input_0, input_1)
}
<
  domain: "pkg.onnxscript.torch_lib.common",
  opset_import: ["" : 18]
>
Rank (input) => (return_val)
{
   tmp = Shape (input)
   return_val = Size (tmp)
}
<
  domain: "pkg.onnxscript.torch_lib.common",
  opset_import: ["" : 18]
>
IsScalar (input) => (return_val)
{
   tmp = Shape (input)
   tmp_0 = Size (tmp)
   tmp_1 = Constant <value_int: int = 0> ()
   return_val = Equal (tmp_0, tmp_1)
}
```

[ghstack-poisoned]

Author: justinchuby

onnxscript/backend/onnx_export.py

@@ -258,9 +258,9 @@ def _rename_variable_s(self, name):
         return str(self._rename_variable(name))
 
     def _rename_domain(self, domain: str) -> str:
-        if domain == "":
-            return "opset"
-        return domain.replace(".", "_")
+        if domain in {"", "ai.onnx"}:
+            return "opset"  #  TODO: Need checks to avoid name conflicts.
+        return _cleanup_variable_name(domain)  # type: ignore[return-value]
 
     def _make_opset_name(self, domain, version):
         return f"{self._rename_domain(domain)}{version}"
@@ -552,11 +552,13 @@ def add_line(line: str) -> None:
         add_line(f"    return {return_values}")
         return "\n".join(result)
 
-    def _translate_graph(self, model: onnx.ModelProto, function_name: str) -> str:
+    def _translate_graph(self, model: onnx.ModelProto, function_name: Optional[str]) -> str:
         graph = model.graph
         opsets = {}
         for imported in model.opset_import:
             opsets[imported.domain] = imported.version
+        if function_name is None:
+            function_name = _cleanup_variable_name(graph.name)
 
         result: list[str] = []
 
@@ -593,7 +595,9 @@ def _import_onnx_types(
             return "from onnxscript.onnx_types import " + ", ".join(sorted_types)
         return ""
 
-    def export(self, proto: onnx.ModelProto | onnx.FunctionProto, function_name: str) -> str:
+    def export(
+        self, proto: onnx.ModelProto | onnx.FunctionProto, function_name: Optional[str]
+    ) -> str:
         result: list[str] = []
 
         def add(line: str) -> None:
@@ -612,7 +616,6 @@ def add(line: str) -> None:
             translated_functions.append(self._translate_graph(proto, function_name))
         else:
             assert isinstance(proto, FunctionProto)
-            # TODO: use function_name?
             translated_functions = [self._translate_function(proto)]
 
         # TODO: unique_function_domain_version.add((f.domain, 1))
@@ -655,22 +658,15 @@ def visit_graph(graph: onnx.GraphProto) -> None:
 
 def export2python(
     model_onnx,
-    opset=None,
-    verbose=True,
-    name=None,
-    rename=False,
-    function_name="main",
-    use_operators=False,
+    function_name: Optional[str] = None,
+    rename: bool = False,
+    use_operators: bool = False,
     inline_const: bool = False,
 ):
     """Exports an ONNX model to the *python* syntax.
 
     Args:
         model_onnx: string or ONNX graph
-        opset: opset to export to (None to select the one from the
-            graph)
-        verbose: inserts prints
-        name: to overwrite onnx name
         rename: rename the names to get shorter names
         function_name: main function name
         use_operators: use Python operators.
@@ -694,9 +690,6 @@ def export2python(
         code = export2python(onx)
         print(code)
     """
-    del opset  # unused
-    del verbose  # unused
-    del name  # unused
     if isinstance(model_onnx, str):
         model_onnx = onnx.load(model_onnx)
 

onnxscript/function_libs/torch_lib/ops/fft.py

@@ -95,11 +95,18 @@ def _fftn_onnx(
     # dimension at the beginning to represent the batch dimension.
     transformed = op.Unsqueeze(self, axes=[0])
 
-    for dim_ in dims:
-        if dim_ >= 0:
-            # Add 1 to account for the batch dimension when counting axes from the left
-            dim_ = dim_ + 1
-        transformed = op.DFT(transformed, axis=dim_, inverse=inverse, onesided=onesided)
+    # Add 1 to account for the batch dimension when counting axes from the left
+    new_dims = [dim_ + 1 if dim_ >= 0 else dim_ for dim_ in dims]
+
+    for dim in new_dims[:-1]:
+        transformed = op.DFT(transformed, axis=dim, inverse=inverse, onesided=False)
+
+    # Torch computers one-sided FFT on the last dimension only.
+    if onesided:
+        transformed = op.DFT(transformed, axis=new_dims[-1], inverse=inverse, onesided=True)
+    else:
+        transformed = op.DFT(transformed, axis=new_dims[-1], inverse=inverse, onesided=False)
+
     # Remove the batch dimension
     transformed = op.Squeeze(transformed, axes=[0])
 

onnxscript/function_libs/torch_lib/ops/nn.py

@@ -1404,10 +1404,18 @@ def aten_pad_sequence(
     raise NotImplementedError()
 
 
-def aten_reflection_pad1d(self: TensorType, padding: INT64) -> TensorType:
+@torch_op("aten::reflection_pad1d")
+def aten_reflection_pad1d(self: TFloat, padding: INT64) -> TFloat:
     """reflection_pad1d(Tensor self, SymInt[2] padding) -> Tensor"""
 
-    raise NotImplementedError()
+    # assert len(padding) == 2
+    # Input of padding argument should be [x,y], need change to onnx format [0, x, 0, y]
+    start = op.Slice(padding, [0], [1], axes=[0])
+    end = op.Slice(padding, [1], [2], axes=[0])
+    padding_onnx = op.Concat(
+        op.Constant(value_ints=[0]), start, op.Constant(value_ints=[0]), end, axis=0
+    )
+    return op.Pad(self, padding_onnx, mode="reflect")
 
 
 def aten_reflection_pad1d_backward(

onnxscript/tests/function_libs/torch_lib/extra_opinfo.py

@@ -190,21 +190,20 @@ def sample_inputs_convolution(op_info, device, dtype, requires_grad, **kwargs):
         )
 
 
-def sample_inputs__fft_c2c(self, device, dtype, requires_grad=False, **_):
-    del self  # Unused
+def _prepare_data_for_fft_ops(device, dtype, requires_grad=False):
     # Adapted from https://github.com/pytorch/pytorch/blob/01069ad4be449f376cf88a56d842b8eb50f6e9b6/torch/testing/_internal/opinfo/core.py#L2448C1-L2541C79
     is_fp16_or_chalf = dtype in (torch.complex32, torch.half)
     if not is_fp16_or_chalf:
-        nd_tensor = functools.partial(
+        oned_tensor = functools.partial(
             opinfo_core.make_tensor,
-            (S, S + 1, S + 2),
+            (31,),
             device=device,
             dtype=dtype,
             requires_grad=requires_grad,
         )
-        oned_tensor = functools.partial(
+        nd_tensor = functools.partial(
             opinfo_core.make_tensor,
-            (31,),
+            (S, S + 1, S + 2),
             device=device,
             dtype=dtype,
             requires_grad=requires_grad,
@@ -214,25 +213,32 @@ def sample_inputs__fft_c2c(self, device, dtype, requires_grad=False, **_):
         high = None
         shapes = ((2, 8, 9), (33,))
 
-        nd_tensor = functools.partial(
+        oned_tensor = functools.partial(
             opinfo_core.make_tensor,
-            shapes[0],
+            shapes[1],
             device=device,
             low=low,
             high=high,
             dtype=dtype,
             requires_grad=requires_grad,
         )
-        oned_tensor = functools.partial(
+        nd_tensor = functools.partial(
             opinfo_core.make_tensor,
-            shapes[1],
+            shapes[0],
             device=device,
             low=low,
             high=high,
             dtype=dtype,
             requires_grad=requires_grad,
         )
 
+    return oned_tensor, nd_tensor
+
+
+def sample_inputs__fft_c2c(self, device, dtype, requires_grad=False, **_):
+    del self  # Unused
+    oned_tensor, nd_tensor = _prepare_data_for_fft_ops(device, dtype, requires_grad)
+
     for normalization, forward in itertools.product((0, 1, 2), (True, False)):
         # 1-D
         yield opinfo_core.SampleInput(
@@ -252,6 +258,29 @@ def sample_inputs__fft_c2c(self, device, dtype, requires_grad=False, **_):
             )
 
 
+def sample_inputs__fft_r2c(self, device, dtype, requires_grad=False, **_):
+    del self  # Unused
+    oned_tensor, nd_tensor = _prepare_data_for_fft_ops(device, dtype, requires_grad)
+
+    for normalization, one_sided in itertools.product((0, 1, 2), (True, True)):
+        # 1-D
+        yield opinfo_core.SampleInput(
+            oned_tensor(), dim=(0,), normalization=normalization, onesided=one_sided
+        )
+        # N-D
+        for dim in [
+            (0,),
+            (1,),
+            (2,),
+            (1, 2),
+            (0, 1),
+            (0, 1, 2),
+        ]:
+            yield opinfo_core.SampleInput(
+                nd_tensor(), dim=dim, normalization=normalization, onesided=one_sided
+            )
+
+
 def sample_inputs_layer_norm(op_info, device, dtype, requires_grad, **kwargs):
     del op_info  # unused
     del kwargs
@@ -1336,6 +1365,25 @@ def sample_inputs__native_batch_norm_legit_no_stats(
             )
 
 
+def sample_inputs_reflection_pad1d(op_info, device, dtype, requires_grad, **kwargs):
+    del op_info
+    del kwargs
+
+    cases: tuple = (  # ignore
+        ((2, 3), (1, 2)),
+        ((4, 5), (0, 1)),
+        ((6, 7), (1, 1)),
+        ((8, 9), (1, 0)),
+    )
+
+    make_inp = opinfo_core.partial(
+        torch.testing.make_tensor, device=device, dtype=dtype, requires_grad=requires_grad
+    )
+
+    for shape, pad in cases:
+        yield opinfo_core.SampleInput(make_inp(shape), args=(pad,))
+
+
 # NOTE: How to create an OpInfo:
 # 1. Create a function that generates sample inputs for the op.
 #    This function should yield SampleInputs.
@@ -1358,6 +1406,13 @@ def sample_inputs__native_batch_norm_legit_no_stats(
         sample_inputs_func=sample_inputs__fft_c2c,
         supports_out=False,
     ),
+    opinfo_core.OpInfo(
+        "ops.aten._fft_r2c",
+        aten_name="_fft_r2c",
+        dtypes=common_dtype.floating_types(),
+        sample_inputs_func=sample_inputs__fft_r2c,
+        supports_out=False,
+    ),
     opinfo_core.OpInfo(
         "ops.aten._local_scalar_dense",
         aten_name="_local_scalar_dense",
@@ -1407,6 +1462,13 @@ def sample_inputs__native_batch_norm_legit_no_stats(
         sample_inputs_func=sample_inputs_convolution,
         supports_out=False,
     ),
+    opinfo_core.OpInfo(
+        "ops.aten.reflection_pad1d",
+        aten_name="ops.aten.reflection_pad1d",
+        dtypes=common_dtype.floating_and_complex_types_and(torch.int64, torch.bfloat16),
+        sample_inputs_func=sample_inputs_reflection_pad1d,
+        supports_out=False,
+    ),
     opinfo_core.OpInfo(
         "ops.aten.index.Tensor",
         aten_name="index.Tensor",

onnxscript/tests/function_libs/torch_lib/ops_test_data.py

@@ -465,6 +465,12 @@ def _where_input_wrangler(
         trace_only=True,
         complex=True,
     ),
+    TorchLibOpInfo(
+        "ops.aten._fft_r2c",  # Custom from extra_opinfo
+        fft_ops.aten__fft_r2c,
+        tolerance={torch.float64: (2e-6, 2e-6), torch.float32: (3e-2, 3e-4)},
+        trace_only=True,
+    ),
     TorchLibOpInfo(
         "ops.aten._local_scalar_dense",
         core_ops.aten__local_scalar_dense,
@@ -1189,6 +1195,13 @@ def _where_input_wrangler(
         matcher=lambda sample: "weight" in sample.kwargs,
         reason="this Aten overload doesn't accept weight as kwargs",
     ),
+    TorchLibOpInfo(
+        "ops.aten.reflection_pad1d",
+        nn_ops.aten_reflection_pad1d,
+    ).xfail(
+        dtypes=(torch.int64,),
+        reason="Torch not implement reflection_pad1d for int64.",
+    ),
     TorchLibOpInfo(
         "nn.functional.reflection_pad2d",
         nn_ops.aten_reflection_pad2d,

requirements/lintrunner/requirements.txt

@@ -3,7 +3,7 @@ lintrunner-adapters>=0.8.0
 # RUFF, RUFF-FIX
 ruff==0.1.6
 # MYPY
-mypy==1.7.0
+mypy==1.7.1
 types-PyYAML==6.0.12.11
 # PYLINT
 pylint==2.17.6

tools/onnx2script.py

@@ -0,0 +1,62 @@
+# -------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+# --------------------------------------------------------------------------
+
+"""
+onnx2script.py
+
+This module provides a script to convert ONNX model files to Python scripts using the onnxscript library.
+
+Usage:
+    python onnx2script.py <input_file> [-o output_file] [-v]
+
+Arguments:
+    input_file: The ONNX model file to convert.
+    -o, --output: The output file name. If not provided, the output will be named after the input file with a .py extension.
+    -v, --verbose: Enables verbose mode. This suppresses the use of overloaded operators and inline constants.
+
+Example:
+    python onnx2script.py model.onnx -o model.py -v
+"""
+
+import argparse
+import os
+from typing import Optional
+
+import onnx
+
+import onnxscript
+
+
+def convert2script(
+    input_file_name: str, output_file_name: Optional[str], verbose: bool
+) -> None:
+    model = onnx.load(input_file_name, load_external_data=False)
+    python_code = onnxscript.proto2python(
+        model, use_operators=not verbose, inline_const=not verbose
+    )
+
+    # If output file name is not provided, use the input file name with .py extension
+    if output_file_name is None:
+        base_name = os.path.splitext(input_file_name)[0]  # Remove extension
+        output_file_name = base_name + ".py"
+
+    with open(output_file_name, "w", encoding="utf-8") as f:
+        f.write(python_code)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Convert ONNX model file to onnxscript file")
+    parser.add_argument("input", help="ONNX model file to convert")
+    parser.add_argument("-o", "--output", help="Output file name")
+    parser.add_argument(
+        "-v",
+        "--verbose",
+        action="store_true",
+        help="Verbose mode, suppresses use of overloaded operators and inline constants",
+        default=False,
+    )
+
+    args = parser.parse_args()
+    convert2script(args.input, args.output, args.verbose)