feat: support aten.index_select converter

py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py

@@ -2782,3 +2782,28 @@ def aten_ops_roll(
         args[1],
         args_bounds_check(args, 2, []),
     )
+
+
+@dynamo_tensorrt_converter(torch.ops.aten.index_select.default)
+@enforce_tensor_types(
+    {
+        0: (TRTTensor,),
+        2: (TRTTensor,),
+    }
+)
+def aten_ops_index_select(
+    ctx: ConversionContext,
+    target: Target,
+    args: Tuple[Argument, ...],
+    kwargs: Dict[str, Argument],
+    name: str,
+) -> Union[TRTTensor, Sequence[TRTTensor]]:
+    return impl.index.index_select(
+        ctx,
+        target,
+        SourceIR.ATEN,
+        name,
+        args[0],
+        args[1],
+        args[2],
+    )

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

@@ -12,6 +12,7 @@
     elementwise,
     embedding,
     grid,
+    index,
     linear,
     matmul,
     normalization,

py/torch_tensorrt/dynamo/conversion/impl/index.py

@@ -0,0 +1,24 @@
+from typing import Optional
+
+from torch.fx.node import Target
+from torch_tensorrt.dynamo._SourceIR import SourceIR
+from torch_tensorrt.dynamo.conversion._ConversionContext import ConversionContext
+from torch_tensorrt.fx.converters.converter_utils import set_layer_name
+from torch_tensorrt.fx.types import TRTTensor
+
+
+def index_select(
+    ctx: ConversionContext,
+    target: Target,
+    source_ir: Optional[SourceIR],
+    name: str,
+    input: TRTTensor,
+    dim: int,
+    index: TRTTensor,
+) -> TRTTensor:
+    # The axis parameter specifies the dimension along which to index.
+    gather_layer = ctx.net.add_gather(input, index, axis=dim)
+
+    set_layer_name(gather_layer, target, f"{name}_gather", source_ir)
+
+    return gather_layer.get_output(0)

tests/py/dynamo/conversion/test_index_select_aten.py

@@ -0,0 +1,38 @@
+import torch
+import torch.nn as nn
+from harness import DispatchTestCase
+from parameterized import param, parameterized
+from torch.testing._internal.common_utils import run_tests
+from torch_tensorrt import Input
+
+
+class TestIndexSelectConverter(DispatchTestCase):
+    @parameterized.expand(
+        [
+            ("1d_input", (10,), 0, (1,)),
+            ("2d_input_dim_0", (10, 3), 0, (0, 2)),
+            ("2d_input_dim_1", (5, 10), 1, (1, 2, 3)),
+            ("3d_input_dim_0", (10, 5, 10), 0, (0, 5)),
+            ("3d_input_dim_2", (10, 5, 10), 2, (3, 3, 4)),
+        ]
+    )
+    def test_index_select(self, _, source_shape, dim, indices_val):
+        class TestIndexSelect(torch.nn.Module):
+            def forward(self, source_tensor, indices_tensor):
+                return torch.ops.aten.index_select.default(
+                    source_tensor, dim, indices_tensor
+                )
+
+        input = [
+            torch.randn(*source_shape, dtype=torch.float32),
+            torch.tensor([*indices_val], dtype=torch.int32),
+        ]
+
+        self.run_test(
+            TestIndexSelect(),
+            input,
+        )
+
+
+if __name__ == "__main__":
+    run_tests()