feat(atenlib): arange with overloads (#285)

arange ops

Author: justinchuby

onnxscript/function_libs/torch_aten/ops/core.py

@@ -13,13 +13,14 @@
 
 from typing import Any, Optional, Sequence, Union
 
-from onnxscript import BOOL, DOUBLE, FLOAT, INT64
+from onnxscript import BOOL, DOUBLE, FLOAT, INT16, INT32, INT64
 from onnxscript.function_libs.torch_aten.registration import torch_op
 from onnxscript.function_libs.torch_aten.typing import (
     TFloat,
     TFloatOrBFloat16,
     TInt,
     TReal,
+    TRealUnlessFloat16OrInt8,
     TRealUnlessInt16OrInt8,
     TTensor,
 )
@@ -226,10 +227,65 @@ def aten_any(self: TensorType) -> TensorType:
     raise NotImplementedError()
 
 
-def aten_arange(end: float) -> TensorType:
+@torch_op("aten::arange")
+def aten_arange(end: Union[DOUBLE, FLOAT, INT16, INT32, INT64], dtype: int = -1) -> TensorType:
     # arange(Scalar end, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None) -> Tensor
 
-    raise NotImplementedError()
+    # Cast input to double if dtype is specified, because the input dtype may be e.g. bool
+    # which Range does not support. The output type is ensured because the output
+    # is casted to the specified dtype.
+    if dtype != -1:
+        end = op.Cast(end, to=DOUBLE.dtype)
+
+    result = op.Range(0, end, 1)
+    if dtype != -1:
+        result = op.Cast(result, to=dtype)
+
+    return result
+
+
+@torch_op("aten::arange", overload=True)
+def aten_arange_start(
+    start: TRealUnlessFloat16OrInt8, end: TRealUnlessFloat16OrInt8, dtype: int = -1
+) -> TensorType:
+    # arange.start(Scalar start, Scalar end, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None) -> Tensor
+
+    # Cast input to double if dtype is specified, because the input dtype may be e.g. bool
+    # which Range does not support. The output type is ensured because the output
+    # is casted to the specified dtype.
+    if dtype != -1:
+        start = op.Cast(start, to=DOUBLE.dtype)
+        end = op.Cast(end, to=DOUBLE.dtype)
+
+    result = op.Range(start, end, 1)
+    if dtype != -1:
+        result = op.Cast(result, to=dtype)
+
+    return result
+
+
+@torch_op("aten::arange", overload=True)
+def aten_arange_start_step(
+    start: TRealUnlessFloat16OrInt8,
+    end: TRealUnlessFloat16OrInt8,
+    step: TRealUnlessFloat16OrInt8,
+    dtype: int = -1,
+) -> TensorType:
+    # arange.start_step(Scalar start, Scalar end, Scalar step=1, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None) -> Tensor
+
+    # Cast input to double if dtype is specified, because the input dtype may be e.g. bool
+    # which Range does not support. The output type is ensured because the output
+    # is casted to the specified dtype.
+    if dtype != -1:
+        start = op.Cast(start, to=DOUBLE.dtype)
+        end = op.Cast(end, to=DOUBLE.dtype)
+        step = op.Cast(step, to=DOUBLE.dtype)
+
+    result = op.Range(start, end, step)
+    if dtype != -1:
+        result = op.Cast(result, to=dtype)
+
+    return result
 
 
 def aten_arccos(self: TensorType) -> TensorType:

onnxscript/function_libs/torch_aten/typing.py

@@ -42,7 +42,7 @@
 ]
 _FloatType = Union[FLOAT16, FLOAT, DOUBLE]
 _IntType = Union[INT8, INT16, INT32, INT64]
-_RealType = Union[
+RealType = Union[
     BFLOAT16,
     FLOAT16,
     FLOAT,
@@ -57,7 +57,10 @@
 TFloat = TypeVar("TFloat", bound=_FloatType)
 TFloatOrBFloat16 = TypeVar("TFloatOrBFloat16", bound=Union[FLOAT16, FLOAT, DOUBLE, BFLOAT16])
 TInt = TypeVar("TInt", bound=_IntType)
-TReal = TypeVar("TReal", bound=_RealType)
+TReal = TypeVar("TReal", bound=RealType)
 TRealUnlessInt16OrInt8 = TypeVar(
     "TRealUnlessInt16OrInt8", bound=Union[FLOAT16, FLOAT, DOUBLE, BFLOAT16, INT32, INT64]
 )
+TRealUnlessFloat16OrInt8 = TypeVar(
+    "TRealUnlessFloat16OrInt8", bound=Union[DOUBLE, FLOAT, INT16, INT32, INT64]
+)

onnxscript/test/function_libs/torch_aten/ops_correctness_test.py

@@ -195,6 +195,9 @@ def _upsample_kwargs_wrangler(kwargs: dict[str, Any]) -> dict[str, Any]:
     "addmm": core_ops.aten_addmm,
     "amax": (core_ops.aten_amax, _amax_amin_kwargs_wrangler),
     "amin": (core_ops.aten_amin, _amax_amin_kwargs_wrangler),
+    "arange_start_step": core_ops.aten_arange_start_step,
+    "arange_start": core_ops.aten_arange_start,
+    "arange": core_ops.aten_arange,
     "asin": core_ops.aten_asin,
     "asinh": core_ops.aten_asinh,
     "atan": core_ops.aten_atan,
@@ -289,6 +292,26 @@ def _upsample_kwargs_wrangler(kwargs: dict[str, Any]) -> dict[str, Any]:
 
 
 SKIP_SUBTESTS: tuple[DecorateMeta, ...] = (
+    skip(
+        "arange",
+        matcher=lambda sample: len(sample.args) != 0,
+        reason="arange overload takes single argument",
+    ),
+    skip(
+        "arange",
+        matcher=lambda sample: sample.kwargs.get("end") is not None,
+        reason="arange overload does not support positional 'end' argument",
+    ),
+    skip(
+        "arange_start",
+        matcher=lambda sample: len(sample.args) != 1,
+        reason="arange_start overload takes two arguments (input, start)",
+    ),
+    skip(
+        "arange_start_step",
+        matcher=lambda sample: len(sample.args) != 2,
+        reason="arange_start_step overload takes three arguments (input, start, step)",
+    ),
     skip(
         "div",
         matcher=lambda sample: sample.kwargs.get("rounding_mode") is not None,
@@ -343,6 +366,15 @@ def _upsample_kwargs_wrangler(kwargs: dict[str, Any]) -> dict[str, Any]:
     ),
 )
 
+duplicate_opinfo(
+    OPS_DB,
+    "arange",
+    (
+        "arange_start",
+        "arange_start_step",
+    ),
+)
+
 
 # END OF SECTION TO MODIFY #####################################################