test(atenlib): split mappings to two (#308)

Split the scripted and traced ops and added a test to make sure we
compiled all functions.

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
 
@@ -2704,7 +2704,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
 

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,
@@ -319,7 +319,6 @@ def _topk_input_wrangler(
     "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,
@@ -332,6 +331,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 = (
@@ -522,6 +541,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",