From 32e0925ab9525d057203bf198a85924cd12541fe Mon Sep 17 00:00:00 2001 From: ShawnNotFound Date: Fri, 30 Oct 2026 22:22:28 +0000 Subject: [PATCH] [Fix][Relax][Frontend][Torch] Handle negative indices in the `select` converter Fixes #20855. `select` passed a negative index straight to `relax.op.take`, whose default `fast` mode does no bounds handling, so `x[:, -1]` read memory before the buffer. Count the index from the end of the dimension when its size is static, and use `wrap` mode when it is symbolic. Add regression tests for the ExportedProgram and FX importers. --- .../torch/exported_program_translator.py | 13 ++++- .../tvm/relax/frontend/torch/fx_translator.py | 13 ++++- .../test_frontend_from_exported_program.py | 58 +++++++++++++++++++ tests/python/relax/test_frontend_from_fx.py | 18 ++++++ 4 files changed, 98 insertions(+), 4 deletions(-) diff --git a/python/tvm/relax/frontend/torch/exported_program_translator.py b/python/tvm/relax/frontend/torch/exported_program_translator.py index a521bfa4b7c7..d7c096865413 100644 --- a/python/tvm/relax/frontend/torch/exported_program_translator.py +++ b/python/tvm/relax/frontend/torch/exported_program_translator.py @@ -1126,8 +1126,17 @@ def _narrow(self, node: fx.Node) -> relax.Var: def _select(self, node: fx.Node) -> relax.Var: x = self.env[node.args[0]] dim = node.args[1] - index = relax.const(node.args[2], "int64") - return self.block_builder.emit(relax.op.take(x, index, dim)) + index = node.args[2] + mode = "fast" + if index < 0: + # PyTorch counts a negative index from the end of the dimension. + dim_size = self.shape_of(x)[dim] + if isinstance(dim_size, tvm_tirx.ir.IntImm): + index += int(dim_size) + else: + mode = "wrap" + index = relax.const(index, "int64") + return self.block_builder.emit(relax.op.take(x, index, dim, mode=mode)) def _slice(self, node: fx.Node) -> relax.Var: import sys diff --git a/python/tvm/relax/frontend/torch/fx_translator.py b/python/tvm/relax/frontend/torch/fx_translator.py index 7d31fa3468f9..88fd71fc1102 100644 --- a/python/tvm/relax/frontend/torch/fx_translator.py +++ b/python/tvm/relax/frontend/torch/fx_translator.py @@ -671,8 +671,17 @@ def _numel(self, node: fx.Node) -> relax.Var: def _select(self, node: fx.Node) -> relax.Var: x = self.env[node.args[0]] dim = node.args[1] - index = relax.const(node.args[2], "int64") - return self.block_builder.emit(relax.op.take(x, index, dim)) + index = node.args[2] + mode = "fast" + if index < 0: + # PyTorch counts a negative index from the end of the dimension. + dim_size = self.shape_of(x)[dim] + if isinstance(dim_size, tvm_tirx.ir.IntImm): + index += int(dim_size) + else: + mode = "wrap" + index = relax.const(index, "int64") + return self.block_builder.emit(relax.op.take(x, index, dim, mode=mode)) def _size(self, node: fx.Node) -> relax.Expr: x = self.env[node.args[0]] diff --git a/tests/python/relax/test_frontend_from_exported_program.py b/tests/python/relax/test_frontend_from_exported_program.py index 85fcf8b825ef..2e6bc61443cd 100644 --- a/tests/python/relax/test_frontend_from_exported_program.py +++ b/tests/python/relax/test_frontend_from_exported_program.py @@ -6990,6 +6990,64 @@ def main( verify_model(Select(), example_args, {}, Expected) +def test_select_negative_index(): + class SelectLast(Module): + def forward(self, input): + return input[:, -1] + + @tvm_tirx.script.ir_module + class Expected: + @R.function + def main( + input: R.Tensor((2, 5, 3), dtype="float32"), + ) -> R.Tuple(R.Tensor((2, 3), dtype="float32")): + with R.dataflow(): + lv: R.Tensor((2, 3), dtype="float32") = R.take( + input, R.const(4, "int64"), axis=1, mode="fast" + ) + gv: R.Tuple(R.Tensor((2, 3), dtype="float32")) = (lv,) + R.output(gv) + return gv + + class SelectFirstDynamic(Module): + def forward(self, input): + return torch.select(input, 0, -2) + + s0 = T.dynamic("s0", "int64") + + @I.ir_module + class ExpectedDynamic: + @R.function + def main( + input: R.Tensor((s0, 3), dtype="float32"), + ) -> R.Tuple(R.Tensor((3,), dtype="float32")): + R.func_attr({"tir_var_lower_bound": {"s31": 2}, "tir_var_upper_bound": {"s31": 64}}) + with R.dataflow(): + lv: R.Tensor((3,), dtype="float32") = R.take( + input, R.const(-2, "int64"), axis=0, mode="wrap" + ) + gv: R.Tuple(R.Tensor((3,), dtype="float32")) = (lv,) + R.output(gv) + return gv + + example_args = (torch.randn(2, 5, 3, dtype=torch.float32),) + verify_model(SelectLast(), example_args, {}, Expected) + verify_model_numerically(SelectLast(), example_args) + + batch = torch.export.Dim("batch", min=2, max=64) + example_args = (torch.randn(4, 3, dtype=torch.float32),) + dynamic_shapes = {"input": {0: batch}} + verify_model( + SelectFirstDynamic(), + example_args, + {}, + ExpectedDynamic, + dynamic_shapes=dynamic_shapes, + map_free_vars=True, + ) + verify_model_numerically(SelectFirstDynamic(), example_args, dynamic_shapes=dynamic_shapes) + + def test_unflatten(): class Unflatten(Module): def forward(self, input): diff --git a/tests/python/relax/test_frontend_from_fx.py b/tests/python/relax/test_frontend_from_fx.py index d3c6315c132b..c78681576b1d 100644 --- a/tests/python/relax/test_frontend_from_fx.py +++ b/tests/python/relax/test_frontend_from_fx.py @@ -5821,6 +5821,24 @@ def main( verify_model(Select(), [([5, 3], "float32")], {}, Expected) + class SelectLast(Module): + def forward(self, data): + return torch.select(data, 1, -1) + + @tvm_tirx.script.ir_module + class ExpectedLast: + @R.function + def main( + inp_0: R.Tensor((5, 3), dtype="float32"), + ) -> R.Tensor((5,), dtype="float32"): + with R.dataflow(): + lv: R.Tensor((5,), dtype="float32") = R.take(inp_0, R.const(2, "int64"), axis=1) + gv: R.Tensor((5,), dtype="float32") = lv + R.output(gv) + return gv + + verify_model(SelectLast(), [([5, 3], "float32")], {}, ExpectedLast) + def test_inplace_copy(): class Copy(Module):