[Bug] [Relax] `InternalError: Divide by zero` when compiling a `reshape` that yields a 0-sized dimension

[raindrops-0199]

Expected behavior

tvm.compile should compile a Relax module containing a reshape whose result has a
0-sized dimension (e.g. reshaping an empty intermediate tensor to (-1, 10), which yields
shape (0, 10)). A 0-sized dimension is a legal empty-tensor shape: eager PyTorch,
torch.export, and the from_exported_program frontend all accept this model, and eager
returns a valid (0, 10) tensor.

Actual behavior

tvm.compile aborts at compile time with:

tvm.error.InternalError: Check failed: pb->value != 0 (0 vs. 0) : Divide by zero
  File ".../src/tirx/op/op.cc", line 623, in tvm::floormod(tvm::PrimExpr, tvm::PrimExpr, tvm::Span)
    if (auto ret = arith::TryConstFold<tirx::FloorMod>(a, b)) return ret.value();
  File ".../src/arith/const_fold.h", line 322, in tvm::arith::TryConstFold(...) [Op = tvm::tirx::FloorMod]
    TVM_FFI_ICHECK_NE(pb->value, 0) << "Divide by zero";

Environment

• TVM: 0.25.dev0 (commit 6b4b866d6)
• PyTorch: 2.12.0+cu130
• Python: 3.12
• OS: Ubuntu x86_64
• Target: llvm (CPU)
• Frontend: tvm.relax.frontend.torch.from_exported_program

Steps to reproduce

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


class M(torch.nn.Module):
    def forward(self, x):
        return x.reshape(-1)[: x.numel() // 10 * 10].reshape(-1, 10)


m = M()
x = torch.randn(3)

# eager mode
out = m(x)
print("shape:", tuple(out.shape))   

# compile mode
ep = torch.export.export(m, (x,))
mod = from_exported_program(ep, keep_params_as_input=True, unwrap_unit_return_tuple=True)
mod, _ = relax.frontend.detach_params(mod)
tvm.compile(mod, target=tvm.target.Target("llvm"))   

Triage

• needs-triage

cc @junrushao