feat: support aten.arange.start_step dynamo converter (#2505)

Author: zewenli98

py/torch_tensorrt/dynamo/conversion/ops_evaluators.py

@@ -2,6 +2,7 @@
 import operator
 from typing import Dict, Sequence, Tuple, Union
 
+import numpy as np
 import torch
 from torch.fx.node import Argument, Node, Target
 from torch_tensorrt.dynamo.conversion._ConversionContext import ConversionContext
@@ -35,3 +36,14 @@ def generic_evaluator(
         f"Evaluating {ConverterRegistry.qualified_name_or_str(target)} on object with name: {name}"
     )
     return target(*args)
+
+
+@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 np.arange(*args)

tests/py/dynamo/conversion/test_arange_aten.py

@@ -0,0 +1,34 @@
+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(
+        [
+            (0, 5, 1),
+            (1, 5, 2),
+            (3, 5, 3),
+            (5, 0, -1),
+            (5, 1, -2),
+            (5, 3, -3),
+        ]
+    )
+    def test_arange(self, start, end, step):
+        class Arange(nn.Module):
+            def forward(self, x):
+                return torch.ops.aten.arange.start_step(start, x.shape[0], step)
+
+        inputs = [torch.randn(end, 1)]
+        self.run_test(
+            Arange(),
+            inputs,
+            use_dynamo_tracer=True,
+        )
+
+
+if __name__ == "__main__":
+    run_tests()