(tvm) branch main updated: [Relax][PyTorch] Fix masked_select VM build (#19937)

[tlopex]

This is an automated email from the ASF dual-hosted git repository.

tlopex pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git


The following commit(s) were added to refs/heads/main by this push:
     new 39e0c7e96c [Relax][PyTorch] Fix masked_select VM build (#19937)
39e0c7e96c is described below

commit 39e0c7e96c9eb97ad660b1ae85955da278f373d2
Author: Hangshuai He <152275085+v-ae...@users.noreply.github.com>
AuthorDate: Fri Jul 10 06:41:32 2026 +0800

    [Relax][PyTorch] Fix masked_select VM build (#19937)
    
    This PR fixes the PyTorch ExportedProgram importer lowering for
    `torch.masked_select`.
    
      Previously, `masked_select` lowered to:
    
      - flatten data and mask
      - `nonzero(mask_flat)`
      - `squeeze(axis=[0])`
      - `take(data_flat, indices)`
    
    However, the result of `R.nonzero` only carried rank information. The
    following `R.squeeze` over the dynamic nonzero output could remain
    unhandled during build/VM execution.
    
    This PR inserts a `match_cast` after `R.nonzero` using the exported
    output metadata, preserving the dynamic selected-length dimension before
    `squeeze`.
    
      A numerical regression test is also added to cover:
    
    PyTorch eager -> torch.export -> Relax import -> build -> VM run ->
    output comparison
    
      Testing:
    
    - `python -m pytest -q
    tests/python/relax/test_frontend_from_exported_program.py -k
    'masked_select'`
---
 .../frontend/torch/base_fx_graph_translator.py     | 10 +++++++++
 .../relax/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 66935c1fba..c62fbf2ace 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 @@ class BaseFXGraphImporter(metaclass=abc.ABCMeta):
         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 afd53f1b74..e4200206a6 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 test_masked_select():
             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 test_masked_select():
     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):