siyiweigeHEW opened a new pull request, #20147:
URL: https://github.com/apache/tvm/pull/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`).
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
