[Bug][Relax][ONNX] LayerNormalization replaces optional Mean and InvStdDev outputs with scalar zeros

[Yuhx141]

Version

• Executed revision: e269315c90e3a061c9e1c77b370ce883b1b223f4
• Current main source rechecked: 0587e59a696cb435ad4d328726a07daa8d966f3c
• TVM: 0.26.dev0
• Host: Ubuntu 24.04 x86-64, LLVM CPU

Reproducer

Run the complete script below from a TVM source environment.
It builds a valid opset-17 ONNX LayerNormalization with all three outputs, imports it with from_onnx, and compares
the compiled Relax function with ONNX's reference evaluator.

Observed output before the assertion fails:

TVM Y: shape=(2, 3, 4)
ONNX Y: shape=(2, 3, 4)
TVM Mean: shape=()
0.0
ONNX Mean: shape=(2, 3, 1)
[[[ 1.5] [ 5.5] [ 9.5]] [[13.5] [17.5] [21.5]]]

The first output matches. The second output has the wrong shape and value. The third output is likewise a scalar
zero, while ONNX returns shape (2, 3, 1) with value 0.89442366 in this example.

Expected behavior

When an ONNX graph observes the optional Mean and InvStdDev outputs, the imported Relax function should return
the saved mean and inverse standard deviation with shape
X.shape[:axis] + (1,) * (X.ndim - axis), as required by the ONNX LayerNormalization-17 contract.

Actual behavior and likely cause

The converter always returns placeholders for outputs two and three:

output = relax.op.nn.layer_norm(data, scale, bias, axis, epsilon)
placeholder = relax.const(0, dtype="float32")
return relax.Tuple([output, placeholder, placeholder])

This silently changes the value and shape of a valid model whenever either optional output is observed. The existing
TVM frontend tests construct only the first output, so they do not exercise this contract. The same placeholder code
remains in current main.

Duplicate check

GitHub issue and pull-request searches for ONNX LayerNormalization, Mean, InvStdDev, and the placeholder comment found
no report of these two outputs being replaced by scalar zeros as of 2026-09-30. Issue #18002 concerns a distinct
LegalizeOps shape failure, and issue #19582 concerns a Windows native crash in another LayerNorm path.

Complete reproducer

import numpy as np
import onnx
import tvm
from onnx import TensorProto, helper
from onnx.reference import ReferenceEvaluator
from tvm import relax
from tvm.relax.frontend.onnx import from_onnx


inputs = [
    helper.make_tensor_value_info("x", TensorProto.FLOAT, [2, 3, 4]),
    helper.make_tensor_value_info("scale", TensorProto.FLOAT, [4]),
    helper.make_tensor_value_info("bias", TensorProto.FLOAT, [4]),
]
outputs = [
    helper.make_tensor_value_info("y", TensorProto.FLOAT, [2, 3, 4]),
    helper.make_tensor_value_info("mean", TensorProto.FLOAT, [2, 3, 1]),
    helper.make_tensor_value_info("inv_std_dev", TensorProto.FLOAT, [2, 3, 1]),
]
node = helper.make_node(
    "LayerNormalization",
    [value.name for value in inputs],
    [value.name for value in outputs],
    axis=-1,
)
model = helper.make_model(
    helper.make_graph([node], "layer_norm", inputs, outputs),
    opset_imports=[helper.make_opsetid("", 17)],
)
onnx.checker.check_model(model)

module = from_onnx(model, keep_params_in_input=True)
executable = relax.build(module, target="llvm", relax_pipeline="default", exec_mode="bytecode")
vm = relax.VirtualMachine(executable, tvm.cpu())
feeds = {
    "x": np.arange(24, dtype="float32").reshape(2, 3, 4),
    "scale": np.ones(4, dtype="float32"),
    "bias": np.zeros(4, dtype="float32"),
}
tvm_output = vm["main"](*(tvm.runtime.tensor(feeds[value.name]) for value in inputs))
onnx_output = ReferenceEvaluator(model).run(None, feeds)

actual_outputs = [output.numpy() for output in tvm_output]
for name, actual, expected in zip(("Y", "Mean", "InvStdDev"), actual_outputs, onnx_output):
    print(f"TVM {name}: shape={actual.shape}\n{actual}")
    print(f"ONNX {name}: shape={expected.shape}\n{expected}")

for actual, expected in zip(actual_outputs, onnx_output):
    np.testing.assert_allclose(actual, expected, rtol=1e-5, atol=1e-5)