diff --git a/python/tvm/relax/frontend/torch/base_fx_graph_translator.py b/python/tvm/relax/frontend/torch/base_fx_graph_translator.py index 21544ed20cfa..df1fe1011a73 100644 --- a/python/tvm/relax/frontend/torch/base_fx_graph_translator.py +++ b/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) diff --git a/tests/python/relax/test_frontend_dynamo.py b/tests/python/relax/test_frontend_dynamo.py index 71a8bd8a468c..0f9eabc59f36 100644 --- a/tests/python/relax/test_frontend_dynamo.py +++ b/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", diff --git a/tests/python/relax/test_frontend_from_fx.py b/tests/python/relax/test_frontend_from_fx.py index 571541eebd36..247c8eac3957 100644 --- a/tests/python/relax/test_frontend_from_fx.py +++ b/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