Merge branch 'main' into new_ops_2

Author: justinchuby

onnxscript/function_libs/torch_aten/ops/core.py

@@ -186,15 +186,15 @@ def aten_alpha_dropout(input: TensorType, p: float, train: bool) -> TensorType:
     raise NotImplementedError()
 
 
-# @torch_op("aten::amax")  # FIXME(#249): Uncomment when CI uses onnx 1.13
+@torch_op("aten::amax")
 def aten_amax(self: TReal, dim: INT64, keepdim: int = 0) -> TReal:
     # amax(Tensor self, int[1] dim=[], bool keepdim=False) -> Tensor
 
     # TODO(justinchuby): Make dim optional, keepdim bool
     return op.ReduceMax(self, dim, keepdims=keepdim)
 
 
-# @torch_op("aten::amin")  # FIXME(#249): Uncomment when CI uses onnx 1.13
+@torch_op("aten::amin")
 def aten_amin(self: TReal, dim: INT64, keepdim: int = 0) -> TReal:
     # amin(Tensor self, int[1] dim=[], bool keepdim=False) -> Tensor
 
@@ -2703,7 +2703,7 @@ def aten_logspace(start: float, end: float, steps: int, base: float = 10.0) -> T
     raise NotImplementedError()
 
 
-@torch_op("aten::logsumexp", trace_only=True)  # FIXME(#249): Script when CI uses onnx 1.13
+@torch_op("aten::logsumexp")
 def aten_logsumexp(self: TReal, dim: INT64, keepdim: int = False) -> TReal:
     # logsumexp(Tensor self, int[1] dim, bool keepdim=False) -> Tensor
 
@@ -4408,16 +4408,30 @@ def aten_sinh(self: TFloat) -> TFloat:
     return op.Sinh(self)
 
 
+@torch_op("aten::slice")
 def aten_slice(
-    self: TensorType,
+    self: TTensor,
     dim: int = 0,
     start: Optional[INT64] = None,
     end: Optional[INT64] = None,
     step: INT64 = 1,
-) -> TensorType:
+) -> TTensor:
     # slice.Tensor(Tensor(a) self, int dim=0, SymInt? start=None, SymInt? end=None, SymInt step=1) -> Tensor(a)
 
-    raise NotImplementedError()
+    # TODO: using OptionalHasElement() to check start/end value
+    start = op.Cast(start, to=INT64.dtype)
+    start = op.Reshape(start, op.Constant(value_ints=[-1]))
+
+    end = op.Cast(end, to=INT64.dtype)
+    end = op.Reshape(end, op.Constant(value_ints=[-1]))
+
+    dim = op.Cast(dim, to=INT64.dtype)
+    dim = op.Reshape(dim, op.Constant(value_ints=[-1]))
+
+    step = op.Cast(step, to=INT64.dtype)
+    step = op.Reshape(step, op.Constant(value_ints=[-1]))
+
+    return op.Slice(self, start, end, dim, step)
 
 
 def aten_slice_backward(

onnxscript/test/function_libs/torch_aten/ops_correctness_test.py

@@ -221,7 +221,9 @@ def _topk_input_wrangler(
 
 # Ops to be tested for numerical consistency between onnx and pytorch
 # Find the names of the OpInfos in torch/testing/_internal/common_methods_invocations.py
-OPINFO_FUNCTION_MAPPING: dict[
+
+# Split the scripted and traced ops to make sure we don't forget to script an op
+OPINFO_FUNCTION_MAPPING_SCRIPTED: dict[
     str,
     onnxscript.OnnxFunction
     | Callable[..., Any]
@@ -245,7 +247,6 @@ def _topk_input_wrangler(
     "atan": core_ops.aten_atan,
     "atanh": core_ops.aten_atanh,
     "bmm": core_ops.aten_bmm,
-    "cat": core_ops.aten_cat,
     "ceil": core_ops.aten_ceil,
     "clamp_max": core_ops.aten_clamp_max,
     "clamp_min": core_ops.aten_clamp_min,
@@ -267,18 +268,17 @@ def _topk_input_wrangler(
     "full": (core_ops.aten_full, _full_input_wrangler),
     "full_like": core_ops.aten_full_like,
     "gt": core_ops.aten_gt,
-    "index_select": core_ops.aten_index_select,
     "isinf": core_ops.aten_isinf,
     "log": core_ops.aten_log,
     "log10": core_ops.aten_log10,
     "log1p": core_ops.aten_log1p,
+    "log_softmax": (special_ops.aten_special_log_softmax, _log_softmax_input_wrangler),
     "log2": core_ops.aten_log2,
     "logaddexp": core_ops.aten_logaddexp,
     "logaddexp2": core_ops.aten_logaddexp2,
     "logcumsumexp": core_ops.aten_logcumsumexp,
     "logdet": core_ops.aten_logdet,
     "logsumexp": (core_ops.aten_logsumexp, _logcumsumexp_input_wrangler),
-    "log_softmax": (special_ops.aten_special_log_softmax, _log_softmax_input_wrangler),
     "lt": core_ops.aten_lt,
     "matmul": core_ops.aten_matmul,
     "mm": core_ops.aten_mm,
@@ -315,12 +315,12 @@ def _topk_input_wrangler(
     "sign": core_ops.aten_sign,
     "sin": core_ops.aten_sin,
     "sinh": core_ops.aten_sinh,
+    "slice": core_ops.aten_slice,
     "sqrt": core_ops.aten_sqrt,
     "sub": core_ops.aten_sub,
     "t": core_ops.aten_t,
     "tan": core_ops.aten_tan,
     "tanh": core_ops.aten_tanh,
-    "transpose": core_ops.aten_transpose,
     "topk": (
         core_ops.aten_topk,
         _topk_input_wrangler,
@@ -333,6 +333,26 @@ def _topk_input_wrangler(
     "zeros_like": core_ops.aten_zeros_like,
 }
 
+
+OPINFO_FUNCTION_MAPPING_TRACE_ONLY: dict[
+    str,
+    Callable[..., Any] | tuple[Callable[..., Any], Callable[..., Any]],
+] = {
+    "cat": core_ops.aten_cat,
+    "index_select": core_ops.aten_index_select,
+    "transpose": core_ops.aten_transpose,
+}
+
+OPINFO_FUNCTION_MAPPING: dict[
+    str,
+    onnxscript.OnnxFunction
+    | Callable[..., Any]
+    | tuple[
+        onnxscript.OnnxFunction | Callable[..., Any],
+        Callable[[list[Any], dict[str, Any]], tuple[list[Any], dict[str, Any]]],
+    ],
+] = {**OPINFO_FUNCTION_MAPPING_SCRIPTED, **OPINFO_FUNCTION_MAPPING_TRACE_ONLY}
+
 TESTED_OPS = frozenset(OPINFO_FUNCTION_MAPPING)
 
 EXPECTED_SKIPS_OR_FAILS = (
@@ -420,6 +440,12 @@ def _topk_input_wrangler(
         matcher=lambda sample: "scale_factor" in sample.kwargs,
         reason="fixme: the scale_factor tests",
     ),
+    skip(
+        "slice",
+        # kwargs {dim, start, end, step} is empty, we cannot give the default value
+        matcher=lambda sample: len(sample.kwargs) == 0,
+        reason="start and end must be 1-D array, cannot be optional, due to ort 1.13 does not support yet",
+    ),
 )
 
 duplicate_opinfo(
@@ -523,6 +549,14 @@ def setUp(self) -> None:
         torch.manual_seed(42)
         np.random.seed(42)
 
+    def test_all_script_functions_are_onnx_functions(self):
+        for func_with_wrangler in OPINFO_FUNCTION_MAPPING_SCRIPTED.values():
+            if isinstance(func_with_wrangler, tuple):
+                func = func_with_wrangler[0]
+            else:
+                func = func_with_wrangler
+            self.assertIsInstance(func, onnxscript.OnnxFunction)
+
     @common_device_type.ops(  # type: ignore[misc]
         [info for info in OPS_DB if info.name in TESTED_OPS],
         allowed_dtypes=TESTED_DTYPES,

onnxscript/test/onnx_backend_test.py

@@ -3,6 +3,8 @@
 # Licensed under the MIT License.
 # --------------------------------------------------------------------------
 
+# pylint: disable=too-many-boolean-expressions
+
 import importlib
 import os
 import unittest

pyproject_pylint.toml

@@ -19,9 +19,7 @@ disable = [
   "redefined-builtin",              # TODO: should we avoid redefined-builtin?
   "too-few-public-methods",
   "too-many-arguments",
-  "too-many-boolean-expressions",
   "too-many-branches",
-  "too-many-function-args",
   "too-many-instance-attributes",
   "too-many-lines",
   "too-many-locals",