onnxscript 0.1.0.dev20231128__tar.gz → 0.1.0.dev20231129__tar.gz

{onnxscript-0.1.0.dev20231128/onnxscript.egg-info → onnxscript-0.1.0.dev20231129}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: onnxscript
-Version: 0.1.0.dev20231128
+Version: 0.1.0.dev20231129
 Summary: Naturally author ONNX functions and models using a subset of Python
 Home-page: https://onnxscript.ai/
 Author-email: Microsoft Corporation <onnx@microsoft.com>

{onnxscript-0.1.0.dev20231128 → onnxscript-0.1.0.dev20231129}/onnxscript/backend/onnx_export.py

@@ -258,9 +258,9 @@ class Exporter:
         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 @@ class Exporter:
         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 @@ class Exporter:
             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 @@ class Exporter:
             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 _attribute_param_types(
 
 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-0.1.0.dev20231128 → onnxscript-0.1.0.dev20231129}/onnxscript/function_libs/torch_lib/_flags.py

@@ -39,3 +39,7 @@ EXPERIMENTAL_INITIALIZERS_AS_INPUTS: bool = _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-0.1.0.dev20231128 → onnxscript-0.1.0.dev20231129}/onnxscript/function_libs/torch_lib/ops/core.py

@@ -86,7 +86,7 @@ def aten__log_softmax_half(
     return aten__log_softmax(self, dim, half_to_float)
 
 
-@torch_op("aten::_log_softmax")
+@torch_op("aten::_log_softmax", traceable=True)
 def aten__log_softmax(
     self: TFloatHighPrecision,
     dim: int,
@@ -251,7 +251,7 @@ def aten_addmv(
     return op.Add(op.Mul(self, beta), op.Mul(op.MatMul(mat, vec), alpha))
 
 
-@torch_op("aten::addr")
+@torch_op("aten::addr", traceable=True)
 def aten_addr(
     self: TReal, vec1: TReal, vec2: TReal, beta: float = 1.0, alpha: float = 1.0
 ) -> TReal:
@@ -329,7 +329,7 @@ def aten_align_to(self: TensorType, names: Sequence[str]) -> TensorType:
     raise NotImplementedError()
 
 
-@torch_op("aten::all")
+@torch_op("aten::all", traceable=True)
 def aten_all(self: TTensor) -> BOOL:
     """all(Tensor self) -> Tensor"""
 
@@ -343,7 +343,7 @@ def aten_all(self: TTensor) -> BOOL:
     return result
 
 
-@torch_op("aten::all.dim")
+@torch_op("aten::all.dim", traceable=True)
 def aten_all_dim(self: TTensor, dim: int, keepdim: bool = False) -> BOOL:
     """all.dim(Tensor self, int dim, bool keepdim=False) -> Tensor"""
 
@@ -371,7 +371,7 @@ def aten_all_dims(self: TTensor, dim: Sequence[int] = (), keepdim: bool = False)
     return self
 
 
-@torch_op("aten::all.dims")
+@torch_op("aten::all.dims", traceable=True)
 def aten_all_dims_no_dim(self: TTensor, keepdims: bool) -> BOOL:
     """all.dims(Tensor self, int[]? dim=None, bool keepdim=False) -> Tensor"""
 
@@ -451,7 +451,7 @@ def aten_angle(self: TensorType) -> TensorType:
     raise NotImplementedError()
 
 
-@torch_op("aten::any")
+@torch_op("aten::any", traceable=True)
 def aten_any(self: TTensor) -> BOOL:
     """any(Tensor self) -> Tensor"""
 
@@ -466,7 +466,7 @@ def aten_any(self: TTensor) -> BOOL:
     return result
 
 
-@torch_op("aten::any.dim")
+@torch_op("aten::any.dim", traceable=True)
 def aten_any_dim(self: TTensor, dim: int, keepdim: bool = False) -> BOOL:
     """any.dim(Tensor self, int dim, bool keepdim=False) -> Tensor"""
 
@@ -496,7 +496,7 @@ def aten_any_dims(self: TTensor, dim: Sequence[int] = (), keepdim: bool = False)
     return self
 
 
-@torch_op("aten::any.dims")
+@torch_op("aten::any.dims", traceable=True)
 def aten_any_dims_no_dim(self: TTensor, keepdims: bool) -> BOOL:
     """any.dims(Tensor self, int[1]? dim=None, bool keepdim=False) -> Tensor"""
 
@@ -688,7 +688,7 @@ def aten_arctanh(self: TensorType) -> TensorType:
     raise NotImplementedError()
 
 
-@torch_op("aten::argmax")
+@torch_op("aten::argmax", traceable=True)
 def aten_argmax(self: Union[RealType, UINT8], keepdim: bool = False) -> INT64:
     """argmax(Tensor self, int? dim=None, bool keepdim=False) -> Tensor"""
 
@@ -701,7 +701,7 @@ def aten_argmax(self: Union[RealType, UINT8], keepdim: bool = False) -> INT64:
     return result
 
 
-@torch_op("aten::argmax")
+@torch_op("aten::argmax", traceable=True)
 def aten_argmax_dim(self: Union[RealType, UINT8], dim: int, keepdim: bool = False) -> INT64:
     """argmax(Tensor self, int? dim=None, bool keepdim=False) -> Tensor"""
 
@@ -888,7 +888,7 @@ def aten_atanh(self: TFloat) -> TFloat:
     return op.Atanh(self)
 
 
-@torch_op("aten::atleast_1d")
+@torch_op("aten::atleast_1d", traceable=True)
 def aten_atleast_1d(self: TTensor) -> TTensor:
     """atleast_1d(Tensor self) -> Tensor"""
 
@@ -936,7 +936,7 @@ def aten_atleast_2d_sequence(self: Sequence[TTensor]) -> TTensor:
     return op.SequenceMap(self, body=reshape_to_2d)
 
 
-@torch_op("aten::atleast_3d")
+@torch_op("aten::atleast_3d", traceable=True)
 def aten_atleast_3d(self: TTensor) -> TTensor:
     """atleast_3d(Tensor self) -> Tensor"""
 
@@ -1548,7 +1548,7 @@ def aten_clamp(self: TReal, min: Optional[TReal] = None, max: Optional[TReal] =
     return clamped
 
 
-@torch_op("aten::clamp_max")
+@torch_op("aten::clamp_max", traceable=True)
 def aten_clamp_max(self: TReal, max_: TReal) -> TReal:
     """clamp_max(Tensor self, Tensor max) -> Tensor"""
 
@@ -1567,7 +1567,7 @@ def aten_clamp_max(self: TReal, max_: TReal) -> TReal:
     return result
 
 
-@torch_op("aten::clamp_min")
+@torch_op("aten::clamp_min", traceable=True)
 def aten_clamp_min(self: TReal, min_: TReal) -> TReal:
     """clamp_min(Tensor self, Tensor min) -> Tensor"""
 
@@ -1946,7 +1946,7 @@ def aten_convolution(
     return result
 
 
-@torch_op("aten::convolution", private=True)
+@torch_op("aten::convolution", private=True, traceable=True)
 def _aten_convolution_onnx(
     input: TFloat,
     weight: TFloat,
@@ -2366,7 +2366,7 @@ def aten_cumsum(
     return _aten_cumsum_onnx(cast, dim)
 
 
-@torch_op("aten::cumsum", private=True)
+@torch_op("aten::cumsum", private=True, traceable=True)
 def _aten_cumsum_onnx(
     self: TRealUnlessInt16OrInt8, dim: Union[INT32, INT64]
 ) -> TRealUnlessInt16OrInt8:
@@ -2451,7 +2451,7 @@ def aten_diagonal(self: TReal, offset: int = 0, dim1: int = 0, dim2: int = 1) ->
     return _aten_diagonal_onnx(self, offset, dim1, dim2, perm, axes)
 
 
-@torch_op("aten::diagonal", private=True)
+@torch_op("aten::diagonal", private=True, traceable=True)
 def _aten_diagonal_onnx(
     self: TTensor, offset: int, dim1: int, dim2: int, perm: Sequence[int], axes: Sequence[int]
 ) -> TTensor:
@@ -2727,7 +2727,7 @@ def aten_dot(self: TFloat, tensor: TFloat) -> TFloat:
     return op.MatMul(self, tensor)
 
 
-@torch_op("aten::dropout")
+@torch_op("aten::dropout", traceable=True)
 def aten_dropout(input: TFloat, p: FLOAT, train: BOOL) -> TFloat:
     """dropout(Tensor input, float p, bool train) -> Tensor"""
 
@@ -3062,7 +3062,7 @@ def aten_embedding_dense_backward(
     raise NotImplementedError()
 
 
-@torch_op("aten::embedding_renorm")
+@torch_op("aten::embedding_renorm", traceable=True)
 def aten_embedding_renorm(
     weight: TFloat, indices: INT64, max_norm: float, norm_type: float = 2.0
 ) -> TFloat:
@@ -3541,7 +3541,7 @@ def aten_fused_moving_avg_obs_fake_quant(
     raise NotImplementedError()
 
 
-@torch_op("aten::gather")
+@torch_op("aten::gather", traceable=True)
 def aten_gather(
     self: TReal,
     dim: int,
@@ -4098,7 +4098,7 @@ def aten_index_reduce(
     raise NotImplementedError()
 
 
-@torch_op("aten::index_select")
+@torch_op("aten::index_select", traceable=True)
 def aten_index_select(self: TTensor, dim: int, index: IntType) -> TTensor:
     """index_select(Tensor self, int dim, Tensor index) -> Tensor"""
 
@@ -4582,7 +4582,7 @@ def aten_logaddexp2(self: TFloatOrBFloat16, other: TFloatOrBFloat16) -> TFloatOr
     return op.Div(op.Log(summation), op.Log(two))
 
 
-@torch_op("aten::logcumsumexp")
+@torch_op("aten::logcumsumexp", traceable=True)
 def aten_logcumsumexp(self: TFloatOrBFloat16, dim: int) -> TFloatOrBFloat16:
     """logcumsumexp(Tensor self, int dim) -> Tensor"""
 
@@ -4696,7 +4696,7 @@ def aten_logspace(start: float, end: float, steps: int, base: float = 10.0) -> T
     raise NotImplementedError()
 
 
-@torch_op("aten::logsumexp")
+@torch_op("aten::logsumexp", traceable=True)
 def aten_logsumexp(self: TFloat, dim: INT64, keepdim: int = False) -> TFloat:
     """logsumexp(Tensor self, int[1] dim, bool keepdim=False) -> Tensor"""
 
@@ -4918,7 +4918,7 @@ def aten_max(self: TReal) -> TReal:
     return result
 
 
-@torch_op("aten::max.dim")
+@torch_op("aten::max.dim", traceable=True)
 def aten_max_dim(self: TReal, dim: int, keepdim: bool = False) -> Tuple[TReal, INT64]:
     """max.dim(Tensor self, int dim, bool keepdim=False) -> (Tensor values, Tensor indices)"""
 
@@ -4954,7 +4954,7 @@ def aten_mean(self: TReal) -> TReal:
     return op.Squeeze(result)
 
 
-@torch_op("aten::mean.dim")
+@torch_op("aten::mean.dim", traceable=True)
 def aten_mean_dim(self: TReal, dim: INT64, keepdim: bool = False) -> TReal:
     """mean.dim(Tensor self, int[1]? dim, bool keepdim=False, *, ScalarType? dtype=None) -> Tensor"""
 
@@ -4986,7 +4986,7 @@ def aten_min(self: TReal) -> TReal:
     return op.ReduceMin(self, keepdims=False)
 
 
-@torch_op("aten::min.dim")
+@torch_op("aten::min.dim", traceable=True)
 def aten_min_dim(self: TReal, dim: int, keepdim: bool = False) -> Tuple[TReal, TInt]:
     """min.dim(Tensor self, int dim, bool keepdim=False) -> (Tensor values, Tensor indices)"""
     if IsScalar(self):
@@ -5479,7 +5479,7 @@ def aten_nansum(
     raise NotImplementedError()
 
 
-@torch_op("aten::narrow")
+@torch_op("aten::narrow", traceable=True)
 def aten_narrow(self: TTensor, dim: INT64, start: INT64, length: INT64) -> TTensor:
     """narrow(Tensor(a) self, int dim, SymInt start, SymInt length) -> Tensor(a)"""
 
@@ -6106,7 +6106,7 @@ def aten_norm_except_dim(v: TensorType, pow: int = 2, dim: int = 0) -> TensorTyp
     raise NotImplementedError()
 
 
-@torch_op(("aten::normal", "aten::normal_functional"))
+@torch_op(("aten::normal", "aten::normal_functional"), traceable=True)
 def aten_normal(
     self: TTensor,
     mean: float = 0.0,
@@ -7675,7 +7675,7 @@ def aten_sum_dim_IntList(
     return result
 
 
-@torch_op("aten::sum", private=True)
+@torch_op("aten::sum", private=True, traceable=True)
 def _aten_sum_dim_onnx(self: TReal, dim: INT64, keepdim: bool = False) -> TReal:
     self_is_scalar = IsScalar(self)
     if self_is_scalar:

{onnxscript-0.1.0.dev20231128 → onnxscript-0.1.0.dev20231129}/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-0.1.0.dev20231128 → onnxscript-0.1.0.dev20231129}/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-0.1.0.dev20231128 → onnxscript-0.1.0.dev20231129}/onnxscript/function_libs/torch_lib/registration.py

@@ -99,6 +99,7 @@ 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.
 
@@ -112,6 +113,18 @@ 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 also be traced. This is an **experimental** flag.
+            A function is traceable if it can both be scripted and traced to produce
+            the same result for a given input. Specifically:
+
+            - A function _can_ be tagged with traceable if its if branches (if any)
+                can be statically evaluated.
+            - A function _should_ be tagged with traceable if it contains if branches
+                and/or CastLike nodes so that they can be evaluated away with the
+                EXPERIMENTAL_PREFER_TRACING on.
+            - A function without if branches or CastLike nodes _should not_ be tagged
+                with traceable because inlining will do the same thing.
+            - A function with `@graph` defined for a `Scan` op is not traceable yet.
     """
     if registry is None:
         registry = default_registry
@@ -128,6 +141,7 @@ def torch_op(
         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-0.1.0.dev20231128 → onnxscript-0.1.0.dev20231129}/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,18 +213,18 @@ 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,
@@ -233,6 +232,13 @@ def sample_inputs__fft_c2c(self, device, dtype, requires_grad=False, **_):
             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 @@ OP_DB: List[opinfo_core.OpInfo] = [
         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 @@ OP_DB: List[opinfo_core.OpInfo] = [
         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-0.1.0.dev20231128 → onnxscript-0.1.0.dev20231129}/onnxscript/tests/function_libs/torch_lib/ops_test_data.py

@@ -465,6 +465,12 @@ TESTED_TORCHLIB_OPS: tuple[TorchLibOpInfo, ...] = (
         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 @@ TESTED_TORCHLIB_OPS: tuple[TorchLibOpInfo, ...] = (
         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,

{onnxscript-0.1.0.dev20231128 → onnxscript-0.1.0.dev20231129}/onnxscript/values.py

@@ -479,6 +479,9 @@ class OnnxFunction(Op):
         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",

{onnxscript-0.1.0.dev20231128 → onnxscript-0.1.0.dev20231129/onnxscript.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: onnxscript
-Version: 0.1.0.dev20231128
+Version: 0.1.0.dev20231129
 Summary: Naturally author ONNX functions and models using a subset of Python
 Home-page: https://onnxscript.ai/
 Author-email: Microsoft Corporation <onnx@microsoft.com>