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 1575515972 [Fix][Relax][ONNX] Fold Min/Max/Sum/Mean constants
elementwise (#20119)
1575515972 is described below
commit 15755159720cb5acefb00a152493da58dfd4f944
Author: Aryan Putta <[email protected]>
AuthorDate: Mon Aug 17 17:12:06 2026 -0400
[Fix][Relax][ONNX] Fold Min/Max/Sum/Mean constants elementwise (#20119)
Fixes #20117.
## Problem
`MultiInputBase._impl_v1` folds all-constant operands with:
```python
output = cls.numpy_op(*np_inputs)
```
`numpy_op` is a reduction (`np.min`, `np.max`, `np.sum`, `np.mean`),
whose signature is `op(a, axis=None, ...)`. Passing the operands
positionally binds the second constant to `axis` instead of combining it
with the first.
Importing a model where `Min`, `Max`, `Sum` or `Mean` has only constant
inputs therefore raises:
```
TypeError: only integer scalar arrays can be converted to a scalar index
```
There is a quieter case. When the second operand is a rank-0 integer
that is a valid axis, numpy accepts it and nothing raises. The fold
returns a reduction of the first operand, with the wrong shape and the
wrong values:
```python
a = np.arange(1, 7).reshape(3, 2) # the first constant
b = np.array(0) # a rank-0 constant, a valid axis
np.min(a, b) # -> [1, 2] shape (2,)
# elementwise min ->
[[0,0],[0,0],[0,0]] shape (3, 2)
```
The issue reports `Min` with two rank-1 constants, but the defect is in
the shared base class, so `Max`, `Sum` and `Mean` are affected
identically.
## Fix
Broadcast the operands, stack them on a new leading axis, then reduce
over it. That is exactly what the non-constant path immediately below
already builds with `broadcast_to`, `stack` and `relax_op`, so both
paths now compute one definition.
## Tests
`test_multi_input_all_constant_inputs` covers all four operators through
`check_correctness`, which compares the imported module against
onnxruntime.
`Sum` and `Mean` accept only floating point operands in ONNX, so the
integer cases use `Min` and `Max`. Those cases pass a rank-0 operand
holding a valid axis index (`0` and `1`), since an out-of-range value
would raise and would not reach the silent path.
Every case fails before this change and passes after: the four float
cases raise `TypeError`, and the two integer cases return a wrong
result.
## Verification
Across 444 combinations of the four operators, six shapes including
rank-0 and broadcasting pairs, three dtypes, and operand counts of one,
three and four, the new fold agrees with the broadcast + stack + reduce
path in every case. Over the same set the old fold raised in 425 and
returned a wrong answer in 7.
The six cases added here were each checked against onnxruntime directly,
and the folded values match its output.
---
python/tvm/relax/frontend/onnx/onnx_frontend.py | 10 ++++--
tests/python/relax/test_frontend_onnx.py | 41 +++++++++++++++++++++++++
2 files changed, 49 insertions(+), 2 deletions(-)
diff --git a/python/tvm/relax/frontend/onnx/onnx_frontend.py
b/python/tvm/relax/frontend/onnx/onnx_frontend.py
index 6bbf220dfe..65bd5bfe1a 100644
--- a/python/tvm/relax/frontend/onnx/onnx_frontend.py
+++ b/python/tvm/relax/frontend/onnx/onnx_frontend.py
@@ -2454,8 +2454,14 @@ 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]):
- np_inputs = [inp.data.numpy() for inp in inputs]
- output = cls.numpy_op(*np_inputs) # pylint: disable=not-callable
+ # 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
+ )
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 4adc3a0ab3..32dd4b0bef 100644
--- a/tests/python/relax/test_frontend_onnx.py
+++ b/tests/python/relax/test_frontend_onnx.py
@@ -815,6 +815,47 @@ def test_multi_input_broadcasting():
)
[email protected](
+ "op_name, dtype, shapes, values",
+ [
+ # Two rank-1 constants, the shape reported in apache/tvm#20117.
+ ("Min", TensorProto.FLOAT, [[2], [2]], [[1.0, 2.0], [3.0, 4.0]]),
+ ("Max", TensorProto.FLOAT, [[2], [2]], [[1.0, 2.0], [3.0, 4.0]]),
+ ("Sum", TensorProto.FLOAT, [[2], [2]], [[1.0, 2.0], [3.0, 4.0]]),
+ ("Mean", TensorProto.FLOAT, [[2], [2]], [[1.0, 2.0], [3.0, 4.0]]),
+ # Sum and Mean accept only floating point operands in ONNX, so the
+ # integer cases cover Min and Max. A rank-0 operand holding a valid
+ # axis index is the input that returned a wrong answer rather than
+ # raising, so it needs 0 or 1 here and not an out-of-range value.
+ ("Min", TensorProto.INT64, [[3, 2], []], [[1, 2, 3, 4, 5, 6], [0]]),
+ ("Max", TensorProto.INT64, [[3, 2], []], [[1, 2, 3, 4, 5, 6], [1]]),
+ ],
+)
+def test_multi_input_all_constant_inputs(op_name, dtype, shapes, values):
+ """Folding constant operands must match the broadcast + stack + reduce
path."""
+ nodes, names = [], []
+ for i, (shape, value) in enumerate(zip(shapes, values)):
+ nodes.append(
+ helper.make_node(
+ "Constant",
+ inputs=[],
+ outputs=[f"c{i}"],
+ value=helper.make_tensor(f"c{i}_v", dtype, shape, value),
+ )
+ )
+ names.append(f"c{i}")
+
+ nodes.append(helper.make_node(op_name, names, ["output"]))
+ output_shape = list(np.broadcast_shapes(*[tuple(shape) for shape in
shapes]))
+ graph = helper.make_graph(
+ nodes,
+ f"all_constant_{op_name}",
+ inputs=[],
+ outputs=[helper.make_tensor_value_info("output", dtype, output_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)
