[Bug][relax.frontend.torch] Segfault in relax::Tuple when importing/building a torch.export program that uses batch-wise advanced indexing write (aten.index_put_)

[tinywisdom]

Expected behavior

from_exported_program followed by relax.build segfaults when the exported PyTorch program performs a batch-wise advanced indexing assignment of the form M[:, rows, cols] = x (two integer index tensors on the last two dims + a leading batch slice). The crash happens in tvm::relax::Tuple construction inside FFI. Plain reads (no assignment) or functional rewrites do not crash.

Actual behavior

!!!!!!! TVM FFI encountered a Segfault !!!!!!! 
  ...
  tvm::relax::Tuple::Tuple(tvm::ffi::Array<tvm::RelaxExpr, void>, tvm::Span) [clone .cold]
  ...
Segmentation fault (core dumped)

Environment

• OS: (Ubuntu 22.04.4 LTS (x86_64))
• TVM version: (release v0.21.0)
• Python: (3.10.16)
• LLVM: (17.0.6)

Steps to reproduce

import torch
import torch.nn as nn
from torch.export import export as torch_export
from tvm.relax.frontend.torch import from_exported_program
import tvm
from tvm import relax

class IndexAssign(nn.Module):
    def forward(self, x):                 # x: [B, 10]
        B = x.size(0)
        M = torch.zeros(B, 11, 11)        # could also set device/dtype to match x
        rows = torch.arange(10)
        cols = rows + 1
        M[:, rows, cols] = x              # batch-wise paired index write
        return M

def main():
    m = IndexAssign().eval()
    inp = torch.randn(2, 10)
    ep = torch_export(m, (inp,))          # export succeeds
    mod = from_exported_program(ep)       # import OK or still OK
    target = tvm.target.Target("llvm")
    relax.build(mod, target=target)       # <-- Segfault here (often)
    print("Built successfully (unexpected).")

if __name__ == "__main__":
    main()

Triage

• needs-triage
• bug