Skip to content

[Fix][Relax][ONNX] Fold Min/Max/Sum/Mean constants elementwise - #20119

Open
aryanputta wants to merge 1 commit into
apache:mainfrom
aryanputta:fix/onnx-multiinput-const-fold
Open

[Fix][Relax][ONNX] Fold Min/Max/Sum/Mean constants elementwise#20119
aryanputta wants to merge 1 commit into
apache:mainfrom
aryanputta:fix/onnx-multiinput-const-fold

Conversation

@aryanputta

Copy link
Copy Markdown

Fixes #20117.

Problem

MultiInputBase._impl_v1 folds all-constant operands with:

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:

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.

MultiInputBase folds all-constant operands with cls.numpy_op(*np_inputs).
numpy_op is a reduction, so the second operand is bound to axis rather than
being combined with the first. Importing a model where Min, Max, Sum or Mean
has only constant inputs raises

    TypeError: only integer scalar arrays can be converted to a scalar index

and when the second operand is a rank-0 integer that happens to be a valid
axis, nothing raises at all: the fold returns a reduction of the first operand,
with the wrong shape and the wrong values.

Broadcast the operands and reduce over a new leading axis, which is what the
non-constant path below already builds with broadcast_to, stack and relax_op.

Tests cover all four operators. Sum and Mean take only floating point operands
in ONNX, so the integer cases use Min and Max, with a rank-0 operand holding a
valid axis index to cover the silent case.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[Bug] Mul + ReduceSum crashes in onnx frontend

1 participant