[Fix][Relax] Drop integer-indexed axes in the Torch FX importer

python/tvm/relax/frontend/torch/base_fx_graph_translator.py

@@ -3067,8 +3067,6 @@ def _getitem(self, node: fx.Node) -> relax.Var:
                 return self.block_builder.emit(relax.TupleGetItem(x, node.args[1]))
 
             assert isinstance(x.ty, relax.TensorType)
-            if isinstance(node.args[1], int):
-                return x
             if not isinstance(node.args[1], list | tuple):
                 indices = [node.args[1]]
             else:
@@ -3079,6 +3077,7 @@ def _getitem(self, node: fx.Node) -> relax.Var:
             stride_end = []
             stride = []
             stride_axes = []
+            squeeze_axes = []
             expand_dim = []
             i = 0
             shape = self.shape_of(x)
@@ -3088,10 +3087,12 @@ def _getitem(self, node: fx.Node) -> relax.Var:
                     non_ellipsis_cnt += 1
             for index in indices:
                 if isinstance(index, int):
-                    stride_begin.append(index)
-                    stride_end.append(index + 1)
+                    start = shape[i] + index if index < 0 else index
+                    stride_begin.append(start)
+                    stride_end.append(start + 1)
                     stride.append(1)
                     stride_axes.append(i)
+                    squeeze_axes.append(i)
                     i = i + 1
                 elif isinstance(index, slice):
                     stride_begin.append(0 if index.start is None else index.start)
@@ -3100,7 +3101,7 @@ def _getitem(self, node: fx.Node) -> relax.Var:
                     stride_axes.append(i)
                     i = i + 1
                 elif index is None:
-                    expand_dim.append(len(stride_axes) + len(expand_dim))
+                    expand_dim.append(i - len(squeeze_axes) + len(expand_dim))
                 elif index is Ellipsis:
                     for _ in range(len(shape) - non_ellipsis_cnt):
                         stride_begin.append(0)
@@ -3133,6 +3134,11 @@ def _getitem(self, node: fx.Node) -> relax.Var:
             sliced = self.block_builder.emit(
                 relax.op.strided_slice(taken, stride_axes, stride_begin, stride_end, stride)
             )
+            if squeeze_axes:
+                selected = self.block_builder.emit(relax.op.squeeze(sliced, squeeze_axes))
+                if expand_dim:
+                    return self.block_builder.emit(relax.op.expand_dims(selected, expand_dim))
+                return selected
             sliced_shape = list(self.shape_of(sliced))
             for i in expand_dim:
                 sliced_shape.insert(i, 1)

tests/python/relax/test_frontend_dynamo.py

@@ -593,6 +593,29 @@ def forward(self, input1):
     verify_dynamo_model(Select2(), [([1, 77, 1280], "float32")], {}, Expected2)
 
 
+@pytest.mark.skipif(not tvm_tirx.testing.device_enabled("llvm"), reason="LLVM is required")
+@pytest.mark.parametrize(
+    "index, rhs_shape",
+    [
+        ((Ellipsis, 0), (3, 2)),
+        (1, (2, 4)),
+        ((Ellipsis, -1), (3, 2)),
+        ((1, None, slice(None), -2, None), (1, 2, 1)),
+        ((None, 2, Ellipsis, 0), (1, 2)),
+        ((1, 0, -1), ()),
+        ((Ellipsis, slice(0, 1)), (3, 2, 1)),
+    ],
+)
+def test_getitem_integer_index_cpu(index, rhs_shape):
+    def model(x, rhs):
+        return x[index] + rhs
+
+    x = torch.arange(24, dtype=torch.float32).reshape(3, 2, 4)
+    rhs = torch.full(rhs_shape, 0.25, dtype=torch.float32)
+    compiled = torch.compile(model, backend=relax_dynamo(), fullgraph=True, dynamic=False)
+    torch.testing.assert_close(compiled(x, rhs), model(x, rhs), rtol=0, atol=0)
+
+
 @pytest.mark.skipif(
     version.parse(torch_version) >= version.parse("2.6.0"),
     reason="Need to support dynamic arange in Relax",

tests/python/relax/test_frontend_from_fx.py

@@ -2671,7 +2671,7 @@ def forward(self, x):
     class expected1:
         @R.function
         def main(x: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tensor(
-            (1, 1, 10, 3), dtype="float32"
+            (1, 10, 3), dtype="float32"
         ):
             # block 0
             with R.dataflow():
@@ -2682,8 +2682,8 @@ def main(x: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tensor(
                     end=[1, T.int64(3), T.int64(10), 3],
                     strides=[1, 2, 1, 1],
                 )
-                lv1: R.Tensor((1, 1, 10, 3), dtype="float32") = R.reshape(lv, (1, 1, 10, 3))
-                gv: R.Tensor((1, 1, 10, 3), dtype="float32") = lv1
+                lv1: R.Tensor((1, 10, 3), dtype="float32") = R.squeeze(lv, axis=[0])
+                gv: R.Tensor((1, 10, 3), dtype="float32") = lv1
                 R.output(gv)
             return gv