[Bug][Relax] FoldBatchnormToConv2D crashes for the documented IOHW kernel layout

[Yuhx141]

Description

The IOHW branch of FoldBatchnormToConv2D calls .reshape() on a Relax Call.
As a result, applying the pass to a well-formed, executable graph using the documented
IOHW kernel layout raises AttributeError.

Environment

• TVM 0.26.dev0, executed at e269315c90e3a061c9e1c77b370ce883b1b223f4
• Ubuntu 24.04 x86-64, LLVM 18.1.3 CPU target
• Current upstream main checked at 043b57fbe9b84aba6c0525fc77b624ef8b7bbe2f
• The relevant Python source is byte-identical at both revisions (SHA-256 c591c953ce9de0f8b539149ab466692572ae608978bbb7ca4875fe919f3c588e)
• I did not have a binary built from current main, so the runtime reproduction below is from the executed revision

Minimal reproduction

import numpy as np
import tvm
from tvm import relax
from tvm.script import ir as I
from tvm.script import relax as R

@I.ir_module
class Module:
    @R.function
    def main(
        x: R.Tensor((1, 1, 2, 2), "float32"),
        weight: R.Tensor((1, 2, 1, 1), "float32"),
        gamma: R.Tensor((2,), "float32"),
        beta: R.Tensor((2,), "float32"),
        mean: R.Tensor((2,), "float32"),
        variance: R.Tensor((2,), "float32"),
    ):
        conv = R.nn.conv2d(x, weight, data_layout="NCHW", kernel_layout="IOHW", out_layout="NCHW", out_dtype="float32")
        return R.nn.batch_norm(conv, gamma, beta, mean, variance, axis=1, training=False)[0]

params = {
    "weight": tvm.runtime.tensor(np.array([[[[1.5]], [[-2.0]]]], "float32")),
    "gamma": tvm.runtime.tensor(np.array([2.0, 3.0], "float32")),
    "beta": tvm.runtime.tensor(np.array([5.0, 7.0], "float32")),
    "mean": tvm.runtime.tensor(np.array([1.0, -1.0], "float32")),
    "variance": tvm.runtime.tensor(np.array([4.0, 9.0], "float32")),
}
source = relax.transform.BindParams("main", params)(Module)
assert relax.analysis.check_well_formed(source, check_ty=True)
relax.transform.FoldBatchnormToConv2D()(source)

Expected behavior

The pass should fold IOHW with a Relax reshape operation, or leave this documented layout unchanged when it cannot fold it.

Actual behavior

The well-formed source executes on LLVM, but the pass raises:

AttributeError: 'Call' object has no attribute 'reshape'

Output-channel extents 1 through 6 all reproduced. Changing only the weight layout/shape to OIHW gave six successful, well-formed folds whose VM outputs matched their sources.

Likely cause

The OIHW branch calls relax.op.reshape(wt, ...), while the IOHW branch calls wt.reshape(...). At this point wt is a Relax Call, not an object with a Python reshape method.

Duplicate check

I searched open and closed TVM issues and PRs for this pass, IOHW, and the exception, and did not find the same cause. #17654 is the original implementation PR rather than a fix.

I can send a fix PR and help with follow-up testing if this diagnosis looks right.

Triage

• needs-triage
• type: bug

[Nanmur]

Thanks for the clear reproducer and diagnosis. I opened #20535 with a focused fix for the IOHW path.

The patch replaces the invalid .reshape() method call on the Relax Call with relax.op.reshape, matching the existing OIHW implementation while keeping the IOHW broadcast shape. It also adds a regression test that checks the transformed module is well formed and that BatchNorm is folded successfully.