Split arg{max,min} into arg{max,min}_dim and improve function signatures | feat(atenlib) (#677)

To remove trace_only and allow the dispatcher to select the correct
function, we explicitly split argmax with argmax_dim and fixed the
function return type.

---

One potential problem as below code:
```python
@torch_op("aten::argmax", overload=True)
def aten_argmax_dim(self: TReal, dim: int, keepdim: bool = False) -> TInt:
    self_is_scaler = op.Size(op.Shape(self)) == 0
    if self_is_scaler:
        self = op.Reshape(self, op.Constant(value_ints=[-1]))

    result = op.ArgMax(self, axis=dim, keepdims=keepdim)
    if self_is_scaler:
        result = op.Squeeze(result)

    return result
```
The above code works well. But when it was rewritten to below (reduce
one if/else), the ShapeInference will fail:

```python
@torch_op("aten::argmax", overload=True)
def aten_argmax_dim(self: TReal, dim: int, keepdim: bool = False) -> TInt:
    self_is_scaler = op.Size(op.Shape(self)) == 0
    if self_is_scaler:
        self = op.Reshape(self, op.Constant(value_ints=[-1]))
        result = op.ArgMax(self, axis=dim, keepdims=keepdim)
        result = op.Squeeze(result)
    else:
        result = op.ArgMax(self, axis=dim, keepdims=keepdim)
    return result
```

---------

Co-authored-by: Justin Chu <justinchu@microsoft.com>

Author: xiaowuhu

onnxscript/function_libs/torch_lib/ops/core.py

@@ -14,7 +14,7 @@
 import math
 from typing import Any, Optional, Sequence, Tuple, Union
 
-from onnxscript import BOOL, DOUBLE, FLOAT, INT8, INT16, INT32, INT64, graph
+from onnxscript import BOOL, DOUBLE, FLOAT, INT8, INT16, INT32, INT64, UINT8, graph
 from onnxscript.function_libs.torch_lib.registration import torch_op
 from onnxscript.function_libs.torch_lib.tensor_typing import (
     IntType,
@@ -520,20 +520,21 @@ def aten_arctanh(self: TensorType) -> TensorType:
     raise NotImplementedError()
 
 
-@torch_op("aten::argmax", trace_only=True)
-def aten_argmax(
-    self: TRealOrUInt8, dim: Optional[int] = None, keepdim: bool = False
-) -> TRealOrUInt8:
+@torch_op("aten::argmax")
+def aten_argmax(self: Union[RealType, UINT8], keepdim: bool = False) -> INT64:
     """argmax(Tensor self, int? dim=None, bool keepdim=False) -> Tensor"""
 
-    if dim is None:  # TODO: use OptionalHasElement(dim)
-        self = op.Reshape(self, op.Constant(value_ints=[-1]))
+    self_is_scaler = op.Size(op.Shape(self)) == 0
+    self = op.Reshape(self, op.Constant(value_ints=[-1]))
+    result = op.ArgMax(self, keepdims=keepdim)
+    if self_is_scaler:
+        result = op.Squeeze(result)
 
-    return _aten_argmax_dim(self, dim=dim, keepdim=keepdim)
+    return result
 
 
-@torch_op("aten::argmax", private=True)
-def _aten_argmax_dim(self: TRealOrUInt8, dim: int, keepdim: bool = False) -> TRealOrUInt8:
+@torch_op("aten::argmax")
+def aten_argmax_dim(self: Union[RealType, UINT8], dim: int, keepdim: bool = False) -> INT64:
     """argmax(Tensor self, int? dim=None, bool keepdim=False) -> Tensor"""
 
     self_is_scaler = op.Size(op.Shape(self)) == 0
@@ -547,20 +548,21 @@ def _aten_argmax_dim(self: TRealOrUInt8, dim: int, keepdim: bool = False) -> TRe
     return result
 
 
-@torch_op("aten::argmin", trace_only=True)
-def aten_argmin(
-    self: TRealOrUInt8, dim: Optional[int] = None, keepdim: bool = False
-) -> TRealOrUInt8:
+@torch_op("aten::argmin")
+def aten_argmin(self: Union[RealType, UINT8], keepdim: bool = False) -> INT64:
     """argmin(Tensor self, int? dim=None, bool keepdim=False) -> Tensor"""
 
-    if dim is None:  # TODO: use OptionalHasElement(dim)
-        self = op.Reshape(self, op.Constant(value_ints=[-1]))
+    self_is_scaler = op.Size(op.Shape(self)) == 0
+    self = op.Reshape(self, op.Constant(value_ints=[-1]))
+    result = op.ArgMin(self, keepdims=keepdim)
+    if self_is_scaler:
+        result = op.Squeeze(result)
 
-    return _aten_argmin_dim(self, dim=dim, keepdim=keepdim)
+    return result
 
 
-@torch_op("aten::argmin", private=True)
-def _aten_argmin_dim(self: TRealOrUInt8, dim: int, keepdim: bool = False) -> TRealOrUInt8:
+@torch_op("aten::argmin")
+def aten_argmin_dim(self: Union[RealType, UINT8], dim: int, keepdim: bool = False) -> INT64:
     """argmin(Tensor self, int? dim=None, bool keepdim=False) -> Tensor"""
 
     self_is_scaler = op.Size(op.Shape(self)) == 0

onnxscript/tests/function_libs/torch_lib/ops_test_data.py

@@ -1209,7 +1209,24 @@ def _where_input_wrangler(
         matcher=lambda sample: sample.kwargs.get("end") is not None,
         reason="arange overload does not support positional 'end' argument",
     ),
-    TorchLibOpInfo("argmax", core_ops.aten_argmax, trace_only=True)
+    TorchLibOpInfo("argmax", core_ops.aten_argmax)
+    .skip(
+        matcher=lambda sample: "dim" in sample.kwargs,
+        reason="this overload does not support the 'dim' attribute by design",
+    )
+    .skip(
+        enabled_if=ops_test_common.IS_WINDOWS,
+        reason="fixme: ORT has memory errors. https://github.com/microsoft/onnxruntime/issues/16492",
+    )
+    .xfail(
+        dtypes=(torch.int64,),
+        reason="fixme: ORT did not implement ArgMax for int64. https://github.com/microsoft/onnxruntime/issues/16654",
+    ),
+    TorchLibOpInfo("argmax_dim", core_ops.aten_argmax_dim)
+    .xfail(
+        matcher=lambda sample: "dim" not in sample.kwargs,
+        reason="this overload requires the 'dim' attribute by design",
+    )
     .skip(
         enabled_if=ops_test_common.IS_WINDOWS,
         reason="fixme: ORT has memory errors. https://github.com/microsoft/onnxruntime/issues/16492",
@@ -1218,7 +1235,24 @@ def _where_input_wrangler(
         dtypes=(torch.int64,),
         reason="fixme: ORT did not implement ArgMax for int64. https://github.com/microsoft/onnxruntime/issues/16654",
     ),
-    TorchLibOpInfo("argmin", core_ops.aten_argmin, trace_only=True)
+    TorchLibOpInfo("argmin", core_ops.aten_argmin)
+    .skip(
+        matcher=lambda sample: "dim" in sample.kwargs,
+        reason="this overload does not support the 'dim' attribute by design",
+    )
+    .skip(
+        enabled_if=ops_test_common.IS_WINDOWS,
+        reason="fixme: ORT has memory errors. https://github.com/microsoft/onnxruntime/issues/16492",
+    )
+    .xfail(
+        dtypes=(torch.int64,),
+        reason="fixme: ORT did not implement ArgMin for int64. https://github.com/microsoft/onnxruntime/issues/16654",
+    ),
+    TorchLibOpInfo("argmin_dim", core_ops.aten_argmin_dim)
+    .xfail(
+        matcher=lambda sample: "dim" not in sample.kwargs,
+        reason="this overload requires the 'dim' attribute by design",
+    )
     .skip(
         enabled_if=ops_test_common.IS_WINDOWS,
         reason="fixme: ORT has memory errors. https://github.com/microsoft/onnxruntime/issues/16492",
@@ -1631,6 +1665,8 @@ def _where_input_wrangler(
 ops_test_common.duplicate_opinfo(OPS_DB, "all", ("all_dim",))
 ops_test_common.duplicate_opinfo(OPS_DB, "any", ("any_dim",))
 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",))