[Relax][PyTorch] Fix masked_select VM build

[V-aerus]

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'

[gemini-code-assist]

Code Review

This pull request updates the _masked_select translation in the PyTorch FX graph translator to extract shape information from tensor_meta and apply a match_cast on the indices. It also updates the expected IR in tests and adds a numerical test for masked_select. The review feedback correctly identifies that match_cast should use relax.TensorStructInfo instead of relax.TensorType, as TensorType does not accept a shape list in its constructor.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[gemini-code-assist]

In TVM Relax, match_cast expects a StructInfo (such as TensorStructInfo) rather than a Type (such as TensorType). Additionally, relax.TensorType's constructor takes ndim and dtype rather than a shape list, so passing [1, num_selected] as the first argument is incorrect. Please use relax.TensorStructInfo instead.

Suggested change

	indices = self.block_builder.match_cast(
	indices, relax.TensorType([1, num_selected], "int64")
	)
	indices = self.block_builder.match_cast(
	indices, relax.TensorStructInfo([1, num_selected], "int64")
	)

[V-aerus]

Thanks for the review. I checked the current TVM Relax API and existing code patterns. BlockBuilder.match_cast is used with relax.TensorType(shape, dtype) in the current codebase, for example in tests/python/relax/test_blockbuilder_core.py, and the ONNX frontend has a similar nonzero -> match_cast pattern using relax.TensorType([1, num_nonzero], "int64").

Also, in this TVM version, relax.TensorType accepts a shape and dtype. I don't see relax.TensorStructInfo being used in the current tree. The current code passes the targeted masked_select tests, so I will keep this as TensorType unless maintainers prefer a different API.