[Bug] [Relax] BatchNorm inferred StructInfo disagrees with legalized output dtypes for FP16 inputs

Start by running the supplied reproducer in both training modes and compare the declared and legalized return types. Then trace Relax BatchNorm type inference and legalization, along with the corresponding TOPI computation; done means inferred, legalized, and runtime output dtypes agree for FP16 inputs without breaking the FP32 controls.

Expected behavior

BatchNorm's inferred Relax return types, legalized return types and runtime tensor dtypes should agree. With all five inputs FP16, Relax infers an FP16 tensor for each of the three outputs. Legalization should preserve that function signature.

Actual behavior

LegalizeOps changes the return dtypes even though the input types are unchanged. The runtime results agree with the changed legalized types:

The default CPU compilation pipeline reproduces this without ConvertLayout or convolution. The input module passes check_well_formed. Both training modes reproduced in three fresh Python processes, each executing three legal inputs with positive moving variance.

The script checks the values against an independent NumPy formula before asserting the dtype contract. The final assertion detects the output type change. Direct FP32 controls preserve their inferred FP32 dtypes; explicitly casting the three outputs back to FP16 also preserves the requested FP16 interface.

The drift is already visible after LegalizeOps, before code generation. The corresponding TOPI BatchNorm computation returns the same promoted dtypes. This report concerns type consistency; the numerical values passing the script's checks do not make a changed return signature valid.

Environment

• Apache TVM 0.26.dev0+source.8f328e8 (source snapshot)
• Ubuntu 22.04.5, Linux 6.8.0-52-generic, x86-64
• Python 3.11.16, LLVM 15.0.7
• Target llvm; no GPU required

Steps to reproduce

Save as repro.py:

# Both run three inputs and fail the final dtype assertion:
python repro.py
python repro.py --training

# Controls pass, in either mode:
python repro.py --dtype float32
python repro.py --training --dtype float32
python repro.py --cast-output
python repro.py --training --cast-output

import argparse
import numpy as np
import tvm
from tvm import relax as R

p = argparse.ArgumentParser()
p.add_argument("--training", action="store_true")
p.add_argument("--dtype", choices=["float16", "float32"], default="float16")
p.add_argument("--cast-output", action="store_true")
args = p.parse_args()
params = [R.Var("x", R.TensorType([2, 2, 2, 2], args.dtype))]
params += [R.Var(name, R.TensorType([2], args.dtype))
           for name in ["gamma", "beta", "moving_mean", "moving_var"]]
bb = R.BlockBuilder()
with bb.function("main", params):
    with bb.dataflow():
        result = bb.emit(R.op.nn.batch_norm(
            *params, axis=1, training=args.training, epsilon=1e-5, momentum=0.1))
        if args.cast_output:
            fields = [bb.emit(R.op.astype(R.TupleGetItem(result, i), args.dtype))
                      for i in range(3)]
            result = R.Tuple(fields)
        out = bb.emit_output(result)
    bb.emit_func_output(out)
mod = bb.get()
assert R.analysis.check_well_formed(mod)
legalized = R.transform.LegalizeOps()(mod)
print("declared:", mod["main"].ret_ty, flush=True)
print("legalized:", legalized["main"].ret_ty, flush=True)
vm = R.VirtualMachine(tvm.compile(mod, target="llvm"), tvm.cpu())
observed = []
for seed in [42, 9001, 20261103]:
    rng = np.random.default_rng(seed)
    x = rng.uniform(-0.25, 0.25, (2, 2, 2, 2)).astype(args.dtype)
    arrays = [x, np.ones(2, args.dtype), np.zeros(2, args.dtype),
              np.zeros(2, args.dtype), np.ones(2, args.dtype)]
    actual = [v.numpy() for v in vm["main"](
        *[tvm.runtime.tensor(a, tvm.cpu()) for a in arrays])]
    data = x.astype("float64")
    if args.training:
        mean = data.mean(axis=(0, 2, 3))
        var = ((data - mean[None, :, None, None]) ** 2).mean(axis=(0, 2, 3))
        expected = [(data - mean[None, :, None, None]) /
                    np.sqrt(var[None, :, None, None] + 1e-5),
                    0.1 * mean, 0.9 + 0.1 * var]
    else:
        expected = [data / np.sqrt(1 + 1e-5), np.zeros(2), np.ones(2)]
    for value, reference in zip(actual, expected):
        np.testing.assert_allclose(value, reference, rtol=2e-3, atol=1e-3)
    observed.append([str(v.dtype) for v in actual])
print("three legal inputs checked; runtime dtypes:", observed, flush=True)
assert all(dtype == args.dtype for row in observed for dtype in row), (
    f"output dtype changed from {args.dtype} to {observed[0]}"
)

Triage

• needs-triage
• type: bug