[Bug] FoldBatchnormToConv2D turns a finite float32 output into NaN

バグは python/tvm/relax/transform/fold_batch_norm_to_conv2d_for_inference.py にあります。まず、wt = scale / sqrt(variance + epsilon) を計算して畳み込みの重みに乗算するロジックを確認してください。float32 のオーバーフローを避けるため、より高い精度（例: float64）で乗算するようスケーリングを調整し、その後 float32 にキャストし直してください。提供された再現スクリプトを実行し、fold 後の出力が有限値のままで、fold 前のベースラインと一致することを確認してください。

Thanks for participating in the TVM community! We use https://discuss.tvm.apache.org/ for any general usage questions and discussions. The issue tracker is used for actionable items such as feature proposals discussion, roadmaps, and bug tracking. You are always welcomed to post on the forum first 😸

Issues that are inactive for a period of time may get closed. We adopt this policy so that we won't lose track of actionable issues that may fall at the bottom of the pile. Feel free to reopen a new one if you feel there is an additional problem that needs attention when an old one gets closed.

Expected behavior

For this valid inference graph, folding BatchNorm into Conv2D should preserve a finite output. The unfused graph returns 50331648.0 on the input below. This report concerns the finite-to-NaN change, not exact bitwise equality under floating-point reassociation.

Actual behavior

With the same input and llvm target, FoldBatchnormToConv2D changes the IR and the result becomes NaN:

• Input after float32 rounding: 9999999827968.0.
• Each original Conv2D product magnitude: 9.999999655936003e+25, below the float32 maximum 3.4028234663852886e+38.
• Before folding: 50331648.0 (finite); after folding: NaN.
• Mathematical result evaluated in float64: about 49999624.14.

The source rewrite computes wt = scale / sqrt(variance + epsilon), then wt_conv = conv_weight * wt and executes Conv2D with wt_conv ([pass source](https://github.com/apache/tvm/blob/main/python/tvm/relax/transform/fold_batch_norm_to_conv2d_for_inference.py)). Here wt is about 1e13. The folded weights remain finite at about 1e26, but their products with the 1e13 input are about 1e39, beyond the float32 limit. Opposite-signed terms can therefore become +Inf and -Inf, and their sum becomes NaN. Before folding, the two Conv2D products are about +1e26 and -1e26 (both representable); they cancel before BatchNorm applies its scale.

Only the folding pass causes this discrepancy in a separate pass isolation check: parameter binding alone and layout conversion alone both preserve the baseline output. This is a numerical regression caused by finite-precision overflow after an algebraically valid transformation. It requires unusually large values; this report does not claim that ordinary inference inputs are affected.

Environment

• TVM 0.26.dev0, locally built from commit 5d4dc8721307cf9677becacf0395c19be130e206.
• Python 3.10.20; Linux x86_64 (WSL2, kernel 6.6.114.1-microsoft-standard-WSL2).
• float32, Relax VM, CPU target llvm; Conv2D NCHW / OIHW; inference BatchNorm with axis=1, center=True, scale=True, training=False.
• Reproduced in two separate fresh processes. The executable reproduction was run on the local commit above; current upstream source was inspected, but a current-main binary was not run.

Steps to reproduce

Run this script with a TVM source build that provides Relax and LLVM:

import numpy as np
import tvm
import tvm_ffi
from tvm import relax
from tvm.script import relax as R

v = np.float32(1e13)
x_data = np.full((1, 2, 1, 1), v, dtype="float32")
weights = np.zeros((1, 2, 3, 3), dtype="float32")
weights[0, 0, 1, 1] = v
weights[0, 1, 1, 1] = -v
params = {
    "weight": weights,
    "scale": np.array([v], dtype="float32"),
    "bias": np.array([v], dtype="float32"),
    "mean": np.array([1.0], dtype="float32"),
    "variance": np.array([1.0], dtype="float32"),
}

@tvm.script.ir_module
class Case:
    @R.function
    def main(x: R.Tensor((1, 2, 1, 1), "float32"),
             weight: R.Tensor((1, 2, 3, 3), "float32"),
             scale: R.Tensor((1,), "float32"),
             bias: R.Tensor((1,), "float32"),
             mean: R.Tensor((1,), "float32"),
             variance: R.Tensor((1,), "float32")) -> R.Tensor((1, 1, 1, 1), "float32"):
        conv = R.nn.conv2d(x, weight, strides=(1, 1), padding=1,
                           dilation=(1, 1), groups=1,
                           data_layout="NCHW", kernel_layout="OIHW")
        bn: R.Tuple(R.Tensor((1, 1, 1, 1), "float32"),
                    R.Tensor((1,), "float32"),
                    R.Tensor((1,), "float32")) = R.nn.batch_norm(
            conv, scale, bias, mean, variance, axis=1, epsilon=1e-5,
            center=True, scale=True, training=False)
        return R.TupleGetItem(bn, 0)

def run(mod):
    exe = relax.build(relax.pipeline.zero_pipeline()(mod), target="llvm")
    vm = relax.VirtualMachine(exe, tvm.cpu())
    return vm["main"](tvm.runtime.tensor(x_data)).numpy().item()

base = relax.transform.ConvertToDataflow(min_size=1)(Case)
bound = relax.transform.BindParams(
    "main", {name: tvm.runtime.tensor(value) for name, value in params.items()}
)(base)
folded = relax.transform.FoldBatchnormToConv2D()(bound)
print("IR changed:", not tvm_ffi.structural_equal(bound, folded))
print("before:", run(bound))
print("after:", run(folded))

Observed in both fresh processes:

IR changed: True
before: 50331648.0
after: nan

The operands and BatchNorm parameters are finite, variance=1 is valid, and the original Conv2D products remain within float32 range. The only difference between the two compiled modules is the explicit FoldBatchnormToConv2D transform; both use the same build pipeline and target.

Related reports [#20530](https://github.com/apache/tvm/issues/20530)–[[#20534](https://github.com/apache/tvm/issues/20534)](https://github.com/apache/tvm/issues/20534) cover disabled scale/center, non-channel axis, training mode, float16 constants, and IOHW layout respectively; the associated open PRs [#20535](https://github.com/apache/tvm/pull/20535)–[[#20537](https://github.com/apache/tvm/pull/20537)](https://github.com/apache/tvm/pull/20537) address those cases. None of these eight items reports this float32 finite-to-NaN overflow case. This is based on the issues and PRs inspected on 2026-10-07, not a claim that every upstream report has been ruled out.

Triage

Please refer to the list of label tags here to find the relevant tags and add them below in a bullet format (example below).

• needs-triage
• type: bug