siyiweigeHEW opened a new issue, #20146:
URL: https://github.com/apache/tvm/issues/20146
### 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`:
```python
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
```python
"""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`:
```python
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
--
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]
