[Bug] LegalizeOps batch_norm uses training-mode computation, ignoring provided running stats

[wuyii8941]

Description

LegalizeOps lowering of relax.nn.batch_norm computes mean and variance from the input data (training mode) instead of using the provided running_mean and running_var parameters (inference mode). This produces incorrect results when batch_norm is used for inference.

When DecomposeOpsForInference is applied before LegalizeOps, the batch_norm is correctly decomposed using the provided running statistics.

Reproducer

import numpy as np
import tvm
from tvm import relax
import tvm.relax.op as R

N, C, HW = 1, 4, 4

bb = relax.BlockBuilder()
x = relax.Var('x', relax.TensorStructInfo((N, C, HW, HW), 'float32'))
gamma = relax.Var('gamma', relax.TensorStructInfo((C,), 'float32'))
beta = relax.Var('beta', relax.TensorStructInfo((C,), 'float32'))
mean = relax.Var('mean', relax.TensorStructInfo((C,), 'float32'))
var = relax.Var('var', relax.TensorStructInfo((C,), 'float32'))
with bb.function('main', [x, gamma, beta, mean, var]):
    with bb.dataflow():
        bn = bb.emit(R.nn.batch_norm(x, gamma, beta, mean, var, axis=1, epsilon=1e-5))
        out = bb.emit_output(relax.TupleGetItem(bn, 0))
    bb.emit_func_output(out)
mod = bb.finalize()

x_np = np.ones((N, C, HW, HW), dtype=np.float32)
gamma_np = np.array([1.0, 2.0, 0.5, 3.0], dtype=np.float32)
beta_np = np.array([0.0, 1.0, -1.0, 0.5], dtype=np.float32)
mean_np = np.array([0.0, 0.5, 1.0, -1.0], dtype=np.float32)
var_np = np.array([1.0, 0.25, 4.0, 0.01], dtype=np.float32)
inputs = [x_np, gamma_np, beta_np, mean_np, var_np]

def run(mod, passes, inputs):
    pipeline = tvm.ir.transform.Sequential(passes)
    mod_l = pipeline(mod)
    exe = tvm.relax.build(mod_l, target='llvm')
    vm = tvm.relax.VirtualMachine(exe, device=tvm.cpu())
    tvm_inputs = [tvm.runtime.tensor(x, device=tvm.cpu()) for x in inputs]
    return vm['main'](*tvm_inputs).numpy()

out_legalize = run(mod, [relax.transform.LegalizeOps()], inputs)
out_correct = run(mod, [relax.transform.DecomposeOpsForInference(),
                         relax.transform.LegalizeOps()], inputs)

# Expected: (1.0 - (-1.0)) / sqrt(0.01 + 1e-5) * 3.0 + 0.5 = 60.47
# Channel 3: legalize=0.5 (wrong), correct=60.47
print(f"LegalizeOps:  {out_legalize[0, 3, 0, 0]:.4f}")   # 0.5000 (WRONG)
print(f"Correct:      {out_correct[0, 3, 0, 0]:.4f}")     # 60.4700

Expected behavior

LegalizeOps batch_norm should use the provided running_mean and running_var for inference, producing the same result as DecomposeOpsForInference.

Actual behavior

LegalizeOps computes mean/var from the input tensor (training-mode batch normalization), producing incorrect results. The max error vs ground truth is 59.97.

Environment

• TVM version: 0.24.dev0 (commit 0b0afd8, 2026-04-24)
• Target: llvm (CPU)

[ConvolutedDog]

It seems that this behavior is consistent with the current Relax semantics. relax.nn.batch_norm defaults to training=True:

tvm/python/tvm/relax/op/nn/nn.py

Lines 1619 to 1741 in b971a75

	def batch_norm(
	data: Expr,
	gamma: Expr,
	beta: Expr,
	moving_mean: Expr,
	moving_var: Expr,
	axis: int,
	epsilon: float = 1e-5,
	center: bool = True,
	scale: bool = True,
	momentum: float = 0.1,
	training: bool = True,
	) -> Expr:
	r"""
	Batch normalization layer (Ioffe and Szegedy, 2014).
	Normalizes the input at each batch, i.e. applies a transformation
	that maintains the mean activation close to 0 and the activation
	standard deviation close to 1.
	.. math::
	data\_mean[i] = mean(data[:,i,:,...]) \\
	data\_var[i] = var(data[:,i,:,...])
	Both *mean* and *var* returns a scalar by treating the input as a vector.
	Then compute the normalized output, which has the same shape as input, as following:
	.. math::
	out[:,i,:,...] = \frac{data[:,i,:,...] - data\_mean[i]}{\sqrt{data\_var[i]+\epsilon}}
	* gamma[i] + beta[i]
	Assume the input has size *k* on axis 1, then both ``gamma`` and ``beta``
	have shape *(k,)*.
	Besides the inputs and the outputs, this operator accepts two auxiliary
	states, ``moving_mean`` and ``moving_var``, which are *k*-length
	vectors. They are global statistics for the whole dataset, which are updated by
	.. code:: python
	moving_mean = moving_mean * momentum + data_mean * (1 - momentum)
	moving_var = moving_var * momentum + data_var * (1 - momentum)
	The parameter ``axis`` specifies which axis of the input shape denotes
	the 'channel' (separately normalized groups). The default is 1.
	Specifying -1 sets the channel axis to be the last item in the input shape.
	.. note::
	This operator has two modes:
	- Training mode.
	- Use the mean and var computed from THIS batch to normalize.
	- Update and then return the running mean and running var.
	- Inference mode.
	- Use the running_mean and running_var parameters to normalize.
	- Do not update the running mean and running var. Just return the original value.
	In the legalization stage, this operator will be legalized to the training mode by default.
	You can use tvm.relax.transform.DecomposeOpsForInference to decompose the operator, so it
	executes the inference mode computation. Similarly, use
	tvm.relax.transform.DecomposeOpsForTraining to execute the training mode computation.
	Parameters
	----------
	data : relax.Expr
	The input data to the operator.
	gamma : relax.Expr
	The gamma scale factor.
	beta : relax.Expr
	The beta offset factor.
	moving_mean : relax.Expr
	Running mean of input.
	moving_var : relax.Expr
	Running variance of input.
	axis : int
	The axis along which the normalization is applied.
	epsilon : float
	Small float added to variance to avoid dividing by zero.
	center : bool
	Indicating if the beta offset will be added to the normalized tensor.
	scale : bool
	Indicating if the gamma scale will be multiplied.
	momentum : float
	The value used for the moving_mean and moving_var update.
	training : bool
	A boolean value to indicate whether training or in eval mode. By default.
	relax batch_norm is training mode. To transform it to inference mode,
	can use DecomposeOpsForInference.
	Returns
	-------
	result : relax.Expr
	The computed result.
	"""
	return _ffi_api.batch_norm( # type: ignore
	data,
	gamma,
	beta,
	moving_mean,
	moving_var,
	axis,
	epsilon,
	center,
	scale,
	momentum,
	training,
	)

When training=True, normalization uses statistics computed from the current batch, not the provided moving_mean and moving_var. Normalization will only use the provided moving_mean and moving_var when you set training=False or manually apply the DecomposeOpsForInference pass.

In your program, out_legalize is calculated based on training, while out_correct is calculated based on inference, so they are naturally different.