[Fix][Relax][Frontend][ONNX] Fix Mean/Sum/Min/Max with all-constant inputs

[siyiweigeHEW]

Fixes: #20146

Summary

The Relax ONNX frontend mishandles Mean / Sum / Min / Max nodes whose
inputs are all constants (model initializers): a single constant input
returns a 0-d scalar (e.g. the global mean 3.5 for a (2, 3) input)
instead of the input unchanged, and multiple constant inputs raise
TypeError: only integer scalar arrays can be converted to a scalar index.
onnxruntime returns correct results in both cases.

Root cause

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

output = cls.numpy_op(*np_inputs)

For Mean / Sum / Min / Max, numpy_op is np.mean / np.sum /
np.min / np.max. These numpy reductions reduce their first argument and
interpret the 2nd and later positional arguments as the axis parameter, not as
additional data tensors. So np.mean(x) reduces the whole tensor to a 0-d
scalar, and np.mean(a, b, ...) passes an array into axis → TypeError.

Fix

Mirror the non-constant path: broadcast each constant to the common shape,
stack along a new leading axis, then reduce along it.

input_shapes = [inp.ty.shape for inp in inputs]
target_shape = tuple(
    int(dim)
    for dim in functools.reduce(compute_broadcast_shape, input_shapes)
)
stacked = _np.stack(
    [_np.broadcast_to(x, target_shape) for x in np_inputs], axis=0
)
output = cls.numpy_op(stacked, axis=0)

A single constant input then reduces a (1, *shape) stack along axis 0, which
is the identity — matching the ONNX semantics. (relax.Constant shape elements
are tvm.tir.IntImm, so they are converted to plain ints for np.broadcast_to.)

Validation

Differential test: Relax (build + VirtualMachine) vs onnxruntime (and
onnx.reference) on the same model, comparing output shapes and max |diff|.

Case	onnxruntime	TVM	max|diff|	Result
Mean, single const (2,3)	(2, 3)	(2, 3)	0.00e+00	OK
Mean, single const (5,)	(5,)	(5,)	0.00e+00	OK
Mean, single const scalar ()	()	()	0.00e+00	OK
Mean, two consts (2,3)	(2, 3)	(2, 3)	0.00e+00	OK
Mean, two consts (2,3,4)	(2, 3, 4)	(2, 3, 4)	0.00e+00	OK
Mean, const broadcast (2,3)+(3,)	(2, 3)	(2, 3)	0.00e+00	OK
Mean, 3 consts broadcast (2,3)+(3,)+(2,1)	(2, 3)	(2, 3)	2.38e-07	OK
Sum, two consts (2,3)	(2, 3)	(2, 3)	0.00e+00	OK
Min, consts int64 (2,3)	(2, 3)	(2, 3)	0.00e+00	OK
Max, single const int32 (2,3)	(2, 3)	(2, 3)	0.00e+00	OK
Mean, non-const broadcast (2,3)+(3,)	(2, 3)	(2, 3)	0.00e+00	OK
Mean, mixed const + graph input	(2, 3)	(2, 3)	0.00e+00	OK

The two failing cases from the issue now match onnxruntime exactly:

onnxruntime -> (2, 3) [1.0, 2.0, 3.0, 4.0, 5.0, 6.0]
TVM        -> (2, 3) [1.0, 2.0, 3.0, 4.0, 5.0, 6.0]
[two inputs] onnxruntime -> (2, 3) [6.0, 7.0, 8.0, 9.0, 10.0, 11.0]
[two inputs] TVM        -> (2, 3) [6.0, 7.0, 8.0, 9.0, 10.0, 11.0]

Run:

python results/TVM/deepseek-v4-flash/prove_hum/onnx_Mean/5修复_差分验证.py

Files changed

• python/tvm/relax/frontend/onnx/onnx_frontend.py — MultiInputBase._impl_v1
  constant-fold path: broadcast + stack + reduce along the leading axis instead
  of numpy_op(*np_inputs) (fixes Mean / Sum / Min / Max).