fix(atenlib): merge aten_clamp_min/max tensor and scalar version; repeat (#262)

- Merge `aten_clamp_min`, `aten_clamp_max`'s scalar and tensor versions
by adding logic on the input rank.
- Fix repeat on the empty `repeats` case.

Co-authored-by: Jay Zhang <[email protected]>

Author: justinchuby

onnxscript/function_libs/torch_aten/ops/core.py

@@ -765,35 +765,42 @@ def aten_clamp(self: TensorType, min_=None, max_=None) -> TensorType:
     return clamped
 
 
-@torch_op("aten::clamp_max.Scalar", overload=True)
-def aten_clamp_max_scalar(self, max_):
-    # clamp_max(Tensor self, Scalar max) -> Tensor
-
-    max_ = op.CastLike(max_, self)
-    return op.Clip(self, None, max_)
-
-
-@torch_op("aten::clamp_max.Tensor")
-def aten_clamp_max_tensor(self, max_):
+@torch_op("aten::clamp_max")
+def aten_clamp_max(self, max_):
     # clamp_max(Tensor self, Tensor max) -> Tensor
 
-    return op.Min(self, max_)
-
+    self_size = op.Size(self)
+    max_shape = op.Shape(max_)
+    max_rank = op.Size(max_shape)
+    if self_size == 0:
+        result = op.Expand(self, max_shape)
+    else:
+        if max_rank == 0:
+            max_ = op.CastLike(max_, self)
+            result = op.Clip(self, None, max_)
+        else:
+            result = op.Min(self, max_)
 
-@torch_op("aten::clamp_min.Scalar", overload=True)
-def aten_clamp_min_scalar(self, min_):
-    # clamp_min(Tensor self, Scalar min) -> Tensor
-    # NOTE: min_ is a rank 0 tensor.
-    # TODO(justinchuby): Specify the type constraints.
-    min_ = op.CastLike(min_, self)
-    return op.Clip(self, min_, None)
+    return result
 
 
-@torch_op("aten::clamp_min.Tensor")
-def aten_clamp_min_tensor(self, min_):
+@torch_op("aten::clamp_min")
+def aten_clamp_min(self, min_):
     # clamp_min(Tensor self, Tensor min) -> Tensor
-    # TODO(justinchuby): Specify the type constraints.
-    return op.Max(self, min_)
+
+    self_size = op.Size(self)
+    min_shape = op.Shape(min_)
+    min_rank = op.Size(min_shape)
+    if self_size == 0:
+        result = op.Expand(self, min_shape)
+    else:
+        if min_rank == 0:
+            min_ = op.CastLike(min_, self)
+            result = op.Clip(self, min_, None)
+        else:
+            result = op.Max(self, min_)
+
+    return result
 
 
 def aten_clone(self: TensorType, memory_format: Optional[str] = None) -> TensorType:
@@ -3976,16 +3983,18 @@ def aten_renorm(self: TensorType, p: float, dim: int, maxnorm: float) -> TensorT
 def aten_repeat(self, repeats: INT64):
     # repeat(Tensor self, SymInt[] repeats) -> Tensor
 
-    # FIXME(justinchuby): When repeats.shape == [0]
-
-    # TODO(justinchuby): Make ones_like a function when onnxscript supports it
-    # shape = ones_like(repeats) := {
-    one = op.Constant(value_int=1)
-    repeats_shape = op.Shape(repeats)
-    shape = op.Expand(one, repeats_shape)
-    # }
-    self_expanded = op.Expand(self, shape)  # type: ignore[arg-type]
-    return op.Tile(self_expanded, repeats)
+    if op.Size(repeats) == 0:
+        result = self
+    else:
+        # TODO(justinchuby): Make ones_like a function when onnxscript supports it
+        # shape = ones_like(repeats) := {
+        one = op.Constant(value_int=1)
+        repeats_shape = op.Shape(repeats)
+        shape = op.Expand(one, repeats_shape)
+        # }
+        self_expanded = op.Expand(self, shape)  # type: ignore[arg-type]
+        result = op.Tile(self_expanded, repeats)
+    return result
 
 
 def aten_repeat_interleave(

onnxscript/test/function_libs/torch_aten/ops_correctness_test.py

@@ -79,7 +79,7 @@ def xfail(
     *,
     reason: str,
     dtypes: Optional[Collection[torch.dtype]] = None,
-):
+) -> DecorateMeta:
     """Expects an OpInfo test to fail.
 
     Args:
@@ -104,7 +104,7 @@ def skip(
     reason: str,
     dtypes: Optional[Collection[torch.dtype]] = None,
     matcher: Optional[Callable[[Any], Any]] = None,
-):
+) -> DecorateMeta:
     """Skips an OpInfo test.
 
     Args:
@@ -171,8 +171,8 @@ def wrapped(fn):
     "atanh": core_ops.aten_atanh,
     "bmm": core_ops.aten_bmm,
     "ceil": core_ops.aten_ceil,
-    "clamp_max": core_ops.aten_clamp_max_tensor,
-    "clamp_min": core_ops.aten_clamp_min_tensor,
+    "clamp_max": core_ops.aten_clamp_max,
+    "clamp_min": core_ops.aten_clamp_min,
     "clamp": core_ops.aten_clamp,
     "cos": core_ops.aten_cos,
     "cosh": core_ops.aten_cosh,
@@ -355,23 +355,7 @@ def wrapped(fn):
 )
 
 
-SKIP_SUBTESTS = (
-    skip(
-        "clamp_max",
-        reason="Empty tensor not yet supported",
-        matcher=lambda sample: sample.input.size() == torch.Size([0]),
-    ),
-    skip(
-        "clamp_min",
-        reason="Empty tensor not yet supported",
-        matcher=lambda sample: sample.input.size() == torch.Size([0]),
-    ),
-    skip(
-        "repeat",
-        reason="repeating when input is a scalar and repeats is empty is not supported",
-        matcher=lambda sample: sample.args[0] == (),
-    ),
-)
+SKIP_SUBTESTS: tuple[DecorateMeta, ...] = ()
 OP_WITH_SKIPPED_SUBTESTS = frozenset(meta.op_name for meta in SKIP_SUBTESTS)
 
 # END OF SECTION TO MODIFY #####################################################

requirements-onnx.txt

@@ -1,5 +1,6 @@
 # Requirements for testing with the release version of onnx and onnxruntime
 
 # TODO(#249): Fix tests for onnx 1.13
+numpy==1.21.5
 onnx==1.12
 onnxruntime