[Bug][Relax][TOPI] rms_norm fails when normalization covers every input axis

Read include/tvm/topi/nn/rms_norm.h, especially the reduction and indexing described in the issue, then run the supplied reproducer in an LLVM-enabled TVM environment. The all-axis ONNX case and direct TOPI construction should succeed while the leading-unit-batch control continues to match the NumPy result.

Problem

Normalizing a FLOAT tensor of shape [3, 4] over axes [0, 1], with a matching [3, 4] weight, fails with Tensor dimension mismatch in read ndim = 1, indices.size=0. Both ONNX RMSNormalization(axis=0) and a direct TOPI construction reproduce the same failure. Adding a leading unit batch dimension and reducing axes [1, 2] gives the same mathematical operation and executes correctly after removing that dimension.

ONNX RMSNormalization-23 normalizes from axis through the final dimension. Axis zero is valid. The NumPy oracle evaluates x / sqrt(mean(x*x) + epsilon) * weight; unoptimized ORT and ONNX ReferenceEvaluator agree.

Minimal reproduction

Use an LLVM-enabled TVM environment with onnx==1.22.0, numpy==2.5.4 and onnxruntime==1.29.0. Save the following as repro.py and run:

OMP_NUM_THREADS=1 OPENBLAS_NUM_THREADS=1 MKL_NUM_THREADS=1 python repro.py

No downloaded model, private package or external input file is required. The script constructs and checks the models, saves the generated models/inputs under artifacts/, and asserts the observed failure as well as its controls. It exits successfully only when those checks reproduce the reported behavior.

from pathlib import Path
import hashlib
import numpy as np
import onnx
import onnxruntime as ort
import tvm
from onnx import TensorProto as T, helper as h, numpy_helper as nh
from onnx.reference import ReferenceEvaluator
from tvm import relax
from tvm.relax.frontend.onnx import from_onnx


def make(nodes, inputs, outputs, initializers=(), opset=18):
    model = h.make_model(h.make_graph(nodes, "repro", inputs, outputs, list(initializers)),
                         opset_imports=[h.make_opsetid("", opset)], ir_version=10)
    onnx.checker.check_model(model, full_check=True)
    return model


def sources(model, feeds, expected):
    artifacts = Path("artifacts")
    artifacts.mkdir(exist_ok=True)
    identity = hashlib.sha256(model.SerializeToString()).hexdigest()
    onnx.save(model, artifacts / (identity + ".onnx"))
    np.savez(artifacts / (identity + "-inputs.npz"), **feeds)
    options = ort.SessionOptions()
    options.intra_op_num_threads = options.inter_op_num_threads = 1
    options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_DISABLE_ALL
    actual = ort.InferenceSession(model.SerializeToString(), options,
                                 providers=["CPUExecutionProvider"]).run(None, feeds)
    reference = ReferenceEvaluator(model).run(None, feeds)
    for values in (actual, reference):
        for a, b in zip(values, expected, strict=True):
            np.testing.assert_allclose(a, b, rtol=1e-5, atol=1e-6)


def tvm_run(model, feeds):
    mod = from_onnx(model, keep_params_in_input=False)
    names = [value.name for value in model.graph.input]
    executable = relax.build(mod, target="llvm", relax_pipeline="default")
    result = relax.VirtualMachine(executable, tvm.cpu())["main"](
        *[tvm.runtime.tensor(feeds[name]) for name in names])
    return [result.numpy()] if hasattr(result, "numpy") else [v.numpy() for v in result]


x = np.linspace(0.25, 2.0, 12, dtype="float32").reshape(3, 4)
w = np.linspace(0.5, 1.5, 12, dtype="float32").reshape(3, 4)
expected = x / np.sqrt(np.mean(x * x) + np.float32(1e-5)) * w
for add_batch in (False, True):
    data = x[None, ...] if add_batch else x
    axis = 1 if add_batch else 0
    model = make([h.make_node("RMSNormalization", ["x", "w"], ["y"],
                              axis=axis, epsilon=1e-5)],
                 [h.make_tensor_value_info("x", T.FLOAT, list(data.shape)),
                  h.make_tensor_value_info("w", T.FLOAT, [3, 4])],
                 [h.make_tensor_value_info("y", T.FLOAT, list(data.shape))], opset=23)
    oracle = expected[None, ...] if add_batch else expected
    sources(model, {"x": data, "w": w}, [oracle])
    try:
        actual = tvm_run(model, {"x": data, "w": w})[0]
    except Exception as error:
        assert not add_batch and "Tensor dimension mismatch" in str(error), str(error)
        print("all-axis ONNX", str(error).splitlines()[-1])
    else:
        assert add_batch, "Reported failure did not reproduce"
        np.testing.assert_allclose(actual, oracle, rtol=1e-5, atol=1e-6)
        print("equivalent leading-unit-batch control passed")

# The same failure occurs directly in TOPI, without ONNX conversion.
from tvm import te, topi
vx = te.placeholder((3, 4), name="x", dtype="float32")
vw = te.placeholder((3, 4), name="w", dtype="float32")
try:
    topi.nn.rms_norm(vx, vw, axis=[0, 1], epsilon=1e-5)
except Exception as error:
    assert "Tensor dimension mismatch" in str(error), str(error)
    print("direct TOPI construction", str(error).splitlines()[-1])
else:
    raise AssertionError("Direct TOPI failure did not reproduce")

Key output

all-axis ONNX Check failed: shape.size() == indices.size() (1 vs. 0) : Tensor dimension mismatch in read ndim = 1, indices.size=0
equivalent leading-unit-batch control passed
direct TOPI construction Check failed: shape.size() == indices.size() (1 vs. 0) : Tensor dimension mismatch in read ndim = 1, indices.size=0

Causal analysis

The TOPI rms_norm implementation constructs its squared sum with keepdims=false and atleast1d=true. When every axis is reduced, that result retains rank one. Its subsequent indexing gathers the non-reduced indices, which are empty, and uses them to read the rank-one sum. The reduction shape and indexing must agree in this all-axis case. The direct TOPI probe isolates this independently of ONNX conversion.

Related work and discovery scope

Targeted searches for rms_norm, all-axis/scalar normalization and dimension failures did not find a covering report. This case uses nonempty input data and valid full-axis normalization; it is not a zero-rank input.

The original observation came from a graph-testing workflow. The self-contained script above is a subsequently constructed minimal diagnostic, not a byte-identical replay of a historical generated model. No global novelty claim is made.

Environment

Linux x86_64 CPU, LLVM target; Python 3.12.3, TVM 0.28.dev0, ONNX 1.22.0, NumPy 2.5.4, ORT 1.29.0, tvm-ffi 0.1.14.post2.dev22+gdf463f9bd.

The exact source revision used to build the installed TVM binaries is unknown. The current-source links below are a separate static inspection at 2bf0532998d4dc7879cb342e873608b90c1d4497; that revision was not built or executed for this report. The observed runtime results refer to these loaded binary identities:

• libtvm_compiler.so SHA256 574a573fe994e5eda7352ddd8d685e64e45c7934bdeeaab57a5a272a77792ab1
• libtvm_runtime.so SHA256 4175916fc0ce6266c912629d2255094b94087425c6ab53ffd5e004d4c292c828
• libtvm_runtime_extra.so SHA256 79c93351fca12dd7851964890cea7da65dd2b2e0e64e5c28bfed56c878b70817
• libtvm_ffi.so SHA256 12072d0ccd9d5cb2c6dac80d42f4217ab50dbf795db03197000614a7a770a816

Reproducer SHA256: 8ece9b58b1db26fef13f7a0a10ce1c3be69f38882d84f4fc7cdcb67ee2912dd2. Model and input hashes can be computed from the files saved by the script. No compiler/library rebuild or environment upgrade was used for the reported checks.