From d131b538b8eea9327a0ece745424433c30e6a2ed Mon Sep 17 00:00:00 2001 From: Hehang Shuai <152275085+V-aerus@users.noreply.github.com> Date: Sat, 4 Jul 2026 16:04:36 +0800 Subject: [PATCH] [Relax][PyTorch] Fix masked_select VM build --- .../torch/base_fx_graph_translator.py | 10 ++++++++ .../test_frontend_from_exported_program.py | 24 +++++++++++++++---- 2 files changed, 30 insertions(+), 4 deletions(-) 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 66935c1fbaf1..c62fbf2ace53 100644 --- a/python/tvm/relax/frontend/torch/base_fx_graph_translator.py +++ b/python/tvm/relax/frontend/torch/base_fx_graph_translator.py @@ -2619,6 +2619,16 @@ def _masked_select(self, node: fx.Node) -> relax.Var: data_flat = self.block_builder.emit(relax.op.reshape(data, [-1])) mask_flat = self.block_builder.emit(relax.op.reshape(mask, [-1])) indices = self.block_builder.emit(relax.op.nonzero(mask_flat)) + tensor_meta = node.meta.get("tensor_meta") + if tensor_meta is not None and len(tensor_meta.shape) == 1: + num_selected = tensor_meta.shape[0] + if not isinstance(num_selected, int): + num_selected = tirx.Var(str(num_selected), "int64") + else: + num_selected = tirx.Var(f"{node.name}_num_selected", "int64") + indices = self.block_builder.match_cast( + indices, relax.TensorType([1, num_selected], "int64") + ) indices_1d = self.block_builder.emit(relax.op.squeeze(indices, axis=[0])) result = self.block_builder.emit(relax.op.take(data_flat, indices_1d, axis=0)) diff --git a/tests/python/relax/test_frontend_from_exported_program.py b/tests/python/relax/test_frontend_from_exported_program.py index 1ee88ea846f7..167daf72d34d 100644 --- a/tests/python/relax/test_frontend_from_exported_program.py +++ b/tests/python/relax/test_frontend_from_exported_program.py @@ -6425,16 +6425,19 @@ def main( data: R.Tensor((2, 3), dtype="float32"), mask: R.Tensor((2, 3), dtype="bool") ) -> R.Tuple(R.Tensor(dtype="float32", ndim=1)): R.func_attr({"tir_var_lower_bound": {"u0": 0}, "tir_var_upper_bound": {"u0": 6}}) + u0 = T.int64() with R.dataflow(): lv: R.Tensor((6,), dtype="float32") = R.reshape(data, R.shape([6])) lv1: R.Tensor((6,), dtype="bool") = R.reshape(mask, R.shape([6])) lv2: R.Tensor(dtype="int64", ndim=2) = R.nonzero(lv1) - lv3: R.Tensor(dtype="int64", ndim=1) = R.squeeze(lv2, axis=[0]) - lv4: R.Tensor(dtype="float32", ndim=1) = R.take(lv, lv3, axis=0, mode="fast") - lv5: R.Tensor((), dtype="int64") = R.const(0, "int64") + lv3: R.Tensor((1, u0), dtype="int64") = R.match_cast( + lv2, R.Tensor((1, u0), dtype="int64") + ) + lv4: R.Tensor((u0,), dtype="int64") = R.squeeze(lv3, axis=[0]) + lv5: R.Tensor((u0,), dtype="float32") = R.take(lv, lv4, axis=0, mode="fast") lv6: R.Tensor((), dtype="bool") = R.const(True, "bool") lv7: R.Tensor((), dtype="bool") = R.const(True, "bool") - gv: R.Tuple(R.Tensor(dtype="float32", ndim=1)) = (lv4,) + gv: R.Tuple(R.Tensor((u0,), dtype="float32")) = (lv5,) R.output(gv) return gv @@ -6445,6 +6448,19 @@ def main( verify_model(MaskedSelect(), example_args, {}, Expected) +@pytest.mark.skipif(not tvm.testing.device_enabled("llvm"), reason="llvm not enabled") +def test_masked_select_numerically(): + class MaskedSelect(Module): + def forward(self, data: torch.Tensor, mask: torch.Tensor): + return torch.masked_select(data, mask) + + example_args = ( + torch.tensor([[1, 2, 3], [4, 5, 6]], dtype=torch.float32), + torch.tensor([[True, False, True], [False, True, False]]), + ) + verify_model_numerically(MaskedSelect(), example_args) + + def test_new_ones(): class NewOnes(Module): def forward(self, x):