[Bug][Relax][Frontend][ONNX] `Mean` with constant (initializer) inputs: single input returns a 0-d scalar, multi-input raises TypeError

[siyiweigeHEW]

Expected behavior

ONNX Mean computes an element-wise mean of its variadic inputs (with Numpy-style
multidirectional broadcasting). With a single input it must return that input
unchanged (the mean of one tensor is the tensor itself); with two inputs it must
return (x1 + x2) / 2. Constant (initializer) inputs are valid and common.

onnxruntime returns the (2, 3) input unchanged for the single-input case and
(2, 3) for the two-input case.

Actual behavior

tvm.relax.frontend.onnx.from_onnx mishandles Mean nodes whose inputs are all
constants (model initializers):

• Single constant input (2, 3): returns a 0-d scalar 3.5 (the global mean of
  1..6) instead of the (2, 3) input unchanged.
• Two constant inputs: raises
  TypeError: only integer scalar arrays can be converted to a scalar index.

Root cause

MultiInputBase._impl_v1's constant-fold path at
python/tvm/relax/frontend/onnx/onnx_frontend.py:2456-2459:

if all([isinstance(inp, relax.Constant) for inp in inputs]):
    np_inputs = [inp.data.numpy() for inp in inputs]
    output = cls.numpy_op(*np_inputs)  # pylint: disable=not-callable
    return relax.const(output, output.dtype)

For Mean, numpy_op = _np.mean, so this calls np.mean(*np_inputs). np.mean
interprets the 2nd and later positional arguments as the axis parameter, not as
additional data tensors:

• single input: np.mean(x) reduces the whole tensor → 0-d scalar;
• multiple inputs: np.mean(a, b, ...) passes array(s) into axis → TypeError.

The same constant-fold path is shared by Sum / Min / Max
(np.sum / np.min / np.max have the identical axis-as-2nd-arg signature),
so all of them are affected the same way.

Environment

• OS: Linux
• TVM: v0.24.dev0 (commit 262c6d2e0; also present on main at 7b2ef6ad5f)
• Python: 3.11
• onnx: 1.20.1
• onnxruntime: 1.24.1

Steps to reproduce

"""Repro: ONNX Mean with constant (initializer) inputs — TVM returns a 0-d scalar for
single input and raises TypeError for multiple inputs; onnxruntime handles both."""
import numpy as np
import onnx, onnxruntime
from onnx import helper, TensorProto
import tvm
from tvm import relax
from tvm.relax.frontend.onnx import from_onnx


def build_mean(consts):
    inits = [
        helper.make_tensor(
            f"c{i}", TensorProto.FLOAT, list(c.shape), c.astype(np.float32).flatten().tolist()
        )
        for i, c in enumerate(consts)
    ]
    Y = helper.make_tensor_value_info("Y", TensorProto.FLOAT, None)
    node = helper.make_node("Mean", [f"c{i}" for i in range(len(consts))], ["Y"])
    graph = helper.make_graph([node], "mean", [], [Y], inits)
    model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)])
    model.ir_version = 8
    return model


# ---- Case 1: single constant input (2,3). ONNX Mean must return the input unchanged. ----
const = (np.arange(6) + 1).reshape(2, 3).astype(np.float32)
m1 = build_mean([const])
ref1 = onnxruntime.InferenceSession(m1.SerializeToString()).run(None, {})[0]
print("onnxruntime ->", ref1.shape, ref1.flatten().tolist())

mod1 = from_onnx(m1)
ex1 = relax.build(mod1, target="llvm")
vm1 = relax.VirtualMachine(ex1, tvm.cpu())
out1 = vm1["main"]().numpy()
print("TVM        ->", out1.shape, out1.flatten().tolist() if out1.ndim else [float(out1)])

# ---- Case 2: two constant inputs (2,3). ONNX Mean returns (x1+x2)/2. ----
m2 = build_mean([const, const + 10])
ref2 = onnxruntime.InferenceSession(m2.SerializeToString()).run(None, {})[0]
print("\n[two inputs] onnxruntime ->", ref2.shape, ref2.flatten().tolist())

try:
    mod2 = from_onnx(m2)
    ex2 = relax.build(mod2, target="llvm")
    vm2 = relax.VirtualMachine(ex2, tvm.cpu())
    out2 = vm2["main"]().numpy()
    print("[two inputs] TVM        ->", out2.shape, out2.flatten().tolist())
except Exception as e:
    print("[two inputs] TVM        ->", type(e).__name__, str(e)[:80])

Actual output:

onnxruntime -> (2, 3) [1.0, 2.0, 3.0, 4.0, 5.0, 6.0]
TVM        -> () [3.5]

[two inputs] onnxruntime -> (2, 3) [6.0, 7.0, 8.0, 9.0, 10.0, 11.0]
[two inputs] TVM        -> TypeError only integer scalar arrays can be converted to a scalar index

Additional context

Fix suggestion: the constant-fold path must reduce element-wise across the input
arrays (with broadcasting) instead of calling np.mean(*np_inputs), e.g. for Mean:

output = functools.reduce(np.add, np_inputs) / len(np_inputs)

or simply route the constant case through the same broadcast_to + stack +
relax.op.mean(stacked, axis=0) path used for non-constant inputs.

Triage

• needs-triage
• bug
• relax
• frontend/onnx