[Bug][Relax][PyTorch] Boolean mask indexing fails in from_exported_program: "TypeError: Take op requires the input indices to have integer dtype. However, the given indices dtype is T.bool" #20947

Expected behavior

Translating a PyTorch module that indexes a tensor with a boolean mask (e.g., x[mask]) via from_exported_program should successfully convert the operation to Relax IR. In PyTorch and NumPy semantics, indexing a tensor with a boolean mask selects elements where the mask evaluates to True.

Actual behavior

When translating an ExportedProgram containing boolean tensor indexing, from_exported_program crashes with a TypeError:

Traceback (most recent call last):
  File "repro.py", line 12, in <module>
    mod = from_exported_program(ep)
  File "python/tvm/relax/frontend/torch/exported_program_translator.py", line 2454, in from_exported_program
    return ExportedProgramImporter().from_exported_program(...)
  File "python/tvm/relax/frontend/torch/exported_program_translator.py", line 2289, in from_exported_program
    output_args = self._translate_fx_graph(...)
  File "python/tvm/relax/frontend/torch/exported_program_translator.py", line 1551, in _translate_fx_graph
    self.env[node] = self.convert_map[func_name](node)
  File "python/tvm/relax/frontend/torch/base_fx_graph_translator.py", line 2331, in _index_tensor
    return self.block_builder.emit(relax.op.take(data, index_tensor, axis=axis))
  File "python/tvm/relax/op/index.py", line 76, in take
    ...
TypeError: Take op requires the input indices to have integer dtype. However, the given indices dtype is T.bool

Environment

• TVM Version: 0.28.dev0 (commit cbb1e12fe on main)
• PyTorch Version: 2.5.1+cpu
• Python Version: 3.12.2
• Operating System: Linux x86_64

Steps to reproduce

import torch
from torch.export import export
from tvm.relax.frontend.torch import from_exported_program

class MaskIndexModel(torch.nn.Module):
    def forward(self, x, mask):
        return x[mask]

x = torch.randn(4)
mask = torch.tensor([True, False, True, False])
ep = export(MaskIndexModel(), (x, mask))
mod = from_exported_program(ep)

Root Cause Analysis

In PyTorch, boolean mask indexing generates the ATen node torch.ops.aten.index.Tensor(x, [mask]).

In python/tvm/relax/frontend/torch/base_fx_graph_translator.py:

# Special case: if there's only one non-None index, use take operation
if len(non_none_indices) == 1:
    axis, index_tensor = non_none_indices[0]
    return self.block_builder.emit(relax.op.take(data, index_tensor, axis=axis))

The implementation unconditionally assumes index_tensor contains integer coordinates and forwards it directly to relax.op.take. When index_tensor has a boolean dtype (T.bool), relax.op.take rejects it because take strictly expects integer index types (int32, int64).

For boolean masks, the mask needs to be converted to coordinates (e.g., via relax.op.nonzero or appropriate mask selection) before indexing into data.

Triage

• needs-triage
• relax:frontend