This is an automated email from the ASF dual-hosted git repository.
tlopex pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new 802e1b784a [Fix][Relax][Frontend][ONNX] Fix Mean/Sum/Min/Max with
all-constant inputs (#20147)
802e1b784a is described below
commit 802e1b784aa1bad91b6d24af55de7175834f001e
Author: HuEnwei <[email protected]>
AuthorDate: Wed Aug 26 06:48:02 2026 +0800
[Fix][Relax][Frontend][ONNX] Fix Mean/Sum/Min/Max with all-constant inputs
(#20147)
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
```python
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.
```python
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:
```bash
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`).
---------
Co-authored-by: FFChopon <[email protected]>
Co-authored-by: Claude <[email protected]>
---
python/tvm/relax/frontend/onnx/onnx_frontend.py | 21 +++++---
tests/python/relax/test_frontend_onnx.py | 67 +++++++++++++++++++++++++
2 files changed, 81 insertions(+), 7 deletions(-)
diff --git a/python/tvm/relax/frontend/onnx/onnx_frontend.py
b/python/tvm/relax/frontend/onnx/onnx_frontend.py
index 7e8616f65f..b63518f6c0 100644
--- a/python/tvm/relax/frontend/onnx/onnx_frontend.py
+++ b/python/tvm/relax/frontend/onnx/onnx_frontend.py
@@ -2468,14 +2468,21 @@ class MultiInputBase(OnnxOpConverter):
if cls.numpy_op is None or cls.relax_op is None:
raise NotImplementedError("numpy_op and relax_op must be defined
for MultiInputBase")
if all([isinstance(inp, relax.Constant) for inp in inputs]):
- # numpy_op is a reduction, so the operands cannot be passed
- # positionally: the second constant would be taken as ``axis``.
- # Broadcast and stack first, then reduce over the stack axis, which
- # is what the non-constant path below builds.
- np_inputs = _np.broadcast_arrays(*[inp.data.numpy() for inp in
inputs])
- output = cls.numpy_op( # pylint: disable=not-callable
- _np.stack(np_inputs, axis=0), axis=0
+ np_inputs = [inp.data.numpy() for inp in inputs]
+ # numpy_op (np.mean/np.sum/np.min/np.max) reduces its first arg,
+ # treating any further positional args as `axis`, so calling it as
+ # numpy_op(*np_inputs) is wrong for the variadic ONNX semantics.
+ # Broadcast to the common shape, stack along a new leading axis,
+ # then reduce along it — mirrors the non-constant path below.
+ 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) # pylint:
disable=not-callable
return relax.const(output, output.dtype)
input_shapes = [inp.ty.shape for inp in inputs]
diff --git a/tests/python/relax/test_frontend_onnx.py
b/tests/python/relax/test_frontend_onnx.py
index 09f3a7b01d..39e067d5d7 100644
--- a/tests/python/relax/test_frontend_onnx.py
+++ b/tests/python/relax/test_frontend_onnx.py
@@ -856,6 +856,73 @@ def test_multi_input_all_constant_inputs(op_name, dtype,
shapes, values):
check_correctness(helper.make_model(graph), opset=13)
[email protected]("op_name", ["Min", "Max", "Sum", "Mean"])
[email protected]("rank", [0, 1, 2, 3])
+def test_multi_input_constant_rank_axis_bounds(op_name, rank):
+ """All-constant inputs reduce over the new leading stack axis.
+
+ Stacking the inputs along ``axis=0`` yields a tensor of rank ``rank+1``,
+ so the reduction axis must lie in ``[-(rank+1), rank]``. ``axis=0`` equals
+ the lower bound ``-(rank+1)`` and is therefore valid for every rank, while
+ a positive ``axis=rank+1`` is out of bounds and must raise. The importer
+ relies on this in ``MultiInputBase._impl_v1``, so both the numpy invariant
+ and the end-to-end folding result are asserted here.
+ """
+ shape = (2,) * rank
+ stacked = np.stack([np.ones(shape, np.float32), np.ones(shape,
np.float32)], axis=0)
+ reduce = {"Min": np.min, "Max": np.max, "Sum": np.sum, "Mean":
np.mean}[op_name]
+ assert stacked.ndim == rank + 1
+ for axis in (0, -(rank + 1), rank, -rank):
+ reduce(stacked, axis=axis) # must not raise
+ for axis in (rank + 1, -(rank + 2)):
+ with pytest.raises((IndexError, ValueError)):
+ reduce(stacked, axis=axis)
+
+ # End-to-end: two identical rank-``rank`` constants fold to the identity.
+ const_nodes = []
+ for name in ("c0", "c1"):
+ const_nodes.append(
+ helper.make_node(
+ "Constant",
+ inputs=[],
+ outputs=[name],
+ value=helper.make_tensor(
+ f"{name}_v", TensorProto.FLOAT, list(shape),
np.ones(shape, np.float32).flatten().tolist()
+ ),
+ )
+ )
+ graph = helper.make_graph(
+ const_nodes + [helper.make_node(op_name, ["c0", "c1"], ["output"])],
+ f"const_rank_{op_name}",
+ inputs=[],
+ outputs=[helper.make_tensor_value_info("output", TensorProto.FLOAT,
list(shape))],
+ )
+ check_correctness(helper.make_model(graph), opset=13)
+
+
[email protected]("op_name", ["Min", "Max", "Sum", "Mean"])
[email protected]("shape", [[], [5], [2, 3]])
+def test_multi_input_constant_single_input(op_name, shape):
+ """A single constant input is returned unchanged (shape preserved)."""
+ if shape:
+ vals = np.arange(np.prod(shape), dtype=np.float32).reshape(shape)
+ else:
+ vals = np.array(1.0, np.float32)
+ node = helper.make_node(
+ "Constant",
+ inputs=[],
+ outputs=["c0"],
+ value=helper.make_tensor("c0_v", TensorProto.FLOAT, shape,
vals.flatten().tolist()),
+ )
+ graph = helper.make_graph(
+ [node, helper.make_node(op_name, ["c0"], ["output"])],
+ f"const_single_{op_name}",
+ inputs=[],
+ outputs=[helper.make_tensor_value_info("output", TensorProto.FLOAT,
shape)],
+ )
+ check_correctness(helper.make_model(graph), opset=13)
+
+
@pytest.mark.parametrize("op_name", ["And", "Or", "Xor"])
def test_binary_bool(op_name: str):
verify_binary(op_name, [32, 32], [32, 32], [32, 32],
dtype=TensorProto.BOOL)
