[Bug][Relax] FoldConstant crashes on shape_to_tensor with symbolic shapes (AssertionError) and folds it into a Constant that breaks reshape

[lrcyyds1]

Expected behavior

shape_to_tensor on a tensor with symbolic dims is a core part of the dynamic-reshape idiom (reshape to a runtime shape). FoldConstant should skip folding it when the shape values are not statically known,
and both official build pipelines should compile valid dynamic-shape modules.

Actual behavior

Two failure modes on the current main:

1. shape_of -> shape_to_tensor -> tensor_to_shape -> reshape with a symbolic dim: the cpu_generic pipeline (relax.get_default_pipeline) crashes with an empty AssertionError inside FoldConstant.
2. shape_of -> shape_to_tensor -> reshape (without tensor_to_shape): FoldConstant folds the shape_to_tensor call into a Constant tensor, which reshape then rejects with TypeError: Reshape requires the input new shape to be Shape.

Environment

OS: Linux x86_64
Target: llvm
TVM commit: 2a2b293c02269f4d9f3526c5b03a7548578e78e8 (current main)

Steps to reproduce

import tvm
from tvm import relax
from tvm import tirx as tir
from tvm.relax import transform

bb = relax.BlockBuilder()
x = relax.Var("x", relax.TensorType([tir.Var("m", "int64")], "float32"))
with bb.function("main", params=[x]):
    with bb.dataflow():
        s = bb.emit(relax.op.shape_of(x), "s")
        t = bb.emit(relax.op.shape_to_tensor(s), "t")
        b = bb.emit(relax.op.tensor_to_shape(t), "b")
        y = bb.emit(relax.op.reshape(x, b), "y")
        gv = bb.emit_output(y)
    bb.emit_func_output(gv)
mod = bb.get()

with tvm.target.Target("llvm"):
    transform.FoldConstant()(mod)   # AssertionError

Behavior breakdown:

• cpu_generic pipeline: AssertionError in FoldConstant
• default pipeline (no FoldConstant): builds and runs correctly
• static control (m replaced by literal 8): passes — a symbolic dim is the necessary condition

Root cause

src/relax/transform/fold_constant.cc, the relax.shape_to_tensor special case:

for (size_t i = 0; i < values.size(); i++) {
  PrimExpr val = values[i];
  arr.push_back(val.as<IntImmNode>()->value);                        // unchecked deref
  is_known &= val.ty().MatchesElementType(DLDataTypeCode::kDLInt, 64);
}

When a shape value is a symbolic variable, the dereference happens before the is_known guard can skip it. Moving the IntImm check ahead of the push_back (treating non-IntImm values as !is_known)
looks sufficient for failure mode 1.

For failure mode 2, the folded result is a runtime Constant tensor while reshape's type contract expects a Shape-typed operand, so the fold itself changes the operand kind.

Suggested fix

Check val->IsInstance<IntImmNode>() (or use the checked form of as<>()) before dereferencing, and only fold when all shape values are concrete IntImms. For mode 2, avoid folding shape_to_tensor into a
Constant when any consumer requires a Shape-typed operand.

Related

• Initially reported in [Bug][Relax] Default LLVM pipeline crashes in FoldConstant when reshape target is a runtime Shape #20193 (no minimal repro / root cause there; this issue adds both). Still reproduces on current main as of commit 2a2b293.

This is the canonical dynamic-reshape idiom, so dynamic-shape models (dynamic batch / sequence length) hit it whenever the fusion pipeline runs.