feat(atenlib): add ops (ArgMax, ArgMin, Detach) (#324)

add ops: argmax, argmin, but use trace_only = True mode, due to
OptionalHasElement(dim) cannot work now.

Author: xiaowuhu

onnxscript/function_libs/torch_aten/ops/core.py

@@ -331,20 +331,38 @@ def aten_arctanh(self: TensorType) -> TensorType:
     raise NotImplementedError()
 
 
-def aten_argmax(
-    self: TensorType, dim: Optional[int] = None, keepdim: bool = False
-) -> TensorType:
+@torch_op("aten::argmax", trace_only=True)
+def aten_argmax(self: TReal, dim: Optional[int] = None, keepdim: bool = False) -> TReal:
     # argmax(Tensor self, int? dim=None, bool keepdim=False) -> Tensor
 
-    raise NotImplementedError()
+    self_is_scaler = op.Size(op.Shape(self)) == 0
+    if self_is_scaler:
+        self = op.Reshape(self, op.Constant(value_ints=[-1]))
+    elif dim is None:  # should use OptionalHasElement(dim)
+        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)
 
-def aten_argmin(
-    self: TensorType, dim: Optional[int] = None, keepdim: bool = False
-) -> TensorType:
+    return result
+
+
+@torch_op("aten::argmin", trace_only=True)
+def aten_argmin(self: TReal, dim: Optional[int] = None, keepdim: bool = False) -> TReal:
     # argmin(Tensor self, int? dim=None, bool keepdim=False) -> Tensor
 
-    raise NotImplementedError()
+    self_is_scaler = op.Size(op.Shape(self)) == 0
+    if self_is_scaler:
+        self = op.Reshape(self, op.Constant(value_ints=[-1]))
+    elif dim is None:  # should use OptionalHasElement(dim)
+        self = op.Reshape(self, op.Constant(value_ints=[-1]))
+
+    result = op.ArgMin(self, axis=dim, keepdims=keepdim)
+    if self_is_scaler:
+        result = op.Squeeze(result)
+
+    return result
 
 
 def aten_argsort(self: TensorType, dim: int = -1, descending: bool = False) -> TensorType:
@@ -1383,10 +1401,11 @@ def aten_det(self: TensorType) -> TensorType:
     raise NotImplementedError()
 
 
+@torch_op("aten::detach")
 def aten_detach(self: TensorType) -> TensorType:
     # detach(Tensor(a) self) -> Tensor(a)
 
-    raise NotImplementedError()
+    return op.Identity(self)
 
 
 def aten_detach_copy(self: TensorType) -> TensorType:

onnxscript/test/function_libs/torch_aten/ops_correctness_test.py

@@ -273,6 +273,7 @@ def _topk_input_wrangler(
     "clone": core_ops.aten_clone,
     "cos": core_ops.aten_cos,
     "cosh": core_ops.aten_cosh,
+    # "detach": core_ops.aten_detach,  # detach is not in OP-TEST-DB
     "div": core_ops.aten_div,
     "dot": core_ops.aten_dot,
     "empty": core_ops.aten_empty,
@@ -364,6 +365,8 @@ def _topk_input_wrangler(
     str,
     Callable[..., Any] | tuple[Callable[..., Any], Callable[..., Any]],
 ] = {
+    "argmax": core_ops.aten_argmax,
+    "argmin": core_ops.aten_argmin,
     "cat": core_ops.aten_cat,
     "index_select": core_ops.aten_index_select,
     "transpose": core_ops.aten_transpose,