feat: support aten.arange.start_step dynamo converter

Author: zewenli98

py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py

@@ -2054,3 +2054,22 @@ def aten_ops_addmm(
         beta=kwargs.get("beta", 1),
         alpha=kwargs.get("alpha", 1),
     )
+
+
+@dynamo_tensorrt_converter(torch.ops.aten.arange.start_step)
+def aten_ops_arange_start_step(
+    ctx: ConversionContext,
+    target: Target,
+    args: Tuple[Argument, ...],
+    kwargs: Dict[str, Argument],
+    name: str,
+) -> Union[TRTTensor, Sequence[TRTTensor]]:
+    return impl.sequence.arange(
+        ctx,
+        target,
+        SourceIR.ATEN,
+        name,
+        args[0],
+        args[1],
+        args_bounds_check(args, 2, 1),
+    )

py/torch_tensorrt/dynamo/conversion/impl/__init__.py

@@ -20,6 +20,7 @@
     pool,
     reduce,
     select,
+    sequence,
     shape,
     shuffle,
     slice,

py/torch_tensorrt/dynamo/conversion/impl/sequence.py

@@ -0,0 +1,21 @@
+from typing import Optional, Union
+
+import numpy as np
+from torch.fx.node import Target
+from torch_tensorrt.dynamo._SourceIR import SourceIR
+from torch_tensorrt.dynamo.conversion._ConversionContext import ConversionContext
+from torch_tensorrt.dynamo.conversion.converter_utils import get_trt_tensor
+from torch_tensorrt.fx.types import TRTTensor
+
+
+def arange(
+    ctx: ConversionContext,
+    target: Target,
+    source_ir: Optional[SourceIR],
+    name: str,
+    start: Union[int, float],
+    end: Union[int, float],
+    step: Union[int, float] = 1,
+) -> TRTTensor:
+    values = np.arange(start, end, step)
+    return get_trt_tensor(ctx, values, f"{name}_arange")

tests/py/dynamo/conversion/test_arange_aten.py

@@ -0,0 +1,41 @@
+import torch
+import torch.nn as nn
+from parameterized import parameterized
+from torch.testing._internal.common_utils import run_tests
+
+from .harness import DispatchTestCase
+
+
+class TestArangeConverter(DispatchTestCase):
+    @parameterized.expand(
+        [
+            # int
+            (0, 5, 1),
+            (1, 5, 2),
+            (3, 5, 3),
+            (5, 0, -1),
+            (5, 1, -2),
+            (5, 3, -3),
+            # float and mixed
+            (0.0, 5.0, 1.0),
+            (1.0, 5.0, 2.2),
+            (2, 10, 3.3),
+            (5.0, 0.0, -3),
+            (5.0, 1, -2.0001),
+            (5, 3.0, -3.9999),
+        ]
+    )
+    def test_arange(self, start, end, step):
+        class Arange(nn.Module):
+            def forward(self):
+                return torch.ops.aten.arange.start_step(start, end, step)
+
+        inputs = []
+        self.run_test(
+            Arange(),
+            inputs,
+        )
+
+
+if __name__ == "__main__":
+    run_tests()