Fix negative-axis handling in ExpandDims shape inference
Merged
TI
4mo ago+65+65+01234567890123456789-2-2-0123456789
PR: Fix negative-axis handling in ExpandDims shape inference
Description
The com.microsoft.ExpandDims type/shape-inference function mishandled negative
axis values, which could lead to an out-of-bounds read during graph resolution
(Graph::Resolve). This PR corrects the axis normalization and makes the scalar
axis read robust to both tensor encodings, and adds a regression test that
exercises the shape-inference path.
Summary of Changes
Fix
| File | Change |
|---|---|
onnxruntime/core/graph/contrib_ops/contrib_defs.cc | Normalize a negative axis against the output rank (rank + axis + 1) instead of the off-by-two rank + axis - 1, so the insertion index stays within [0, rank]. Read the scalar axis via the existing ParseScalar helper, which handles both raw_data and int32_data encodings and validates the element count. |
Test
| File | Change |
|---|---|
onnxruntime/test/contrib_ops/expand_dims_test.cc | Add ExpandDimsTest.NegativeAxisConstInitializerShapeInference plus a RunExpandDimsConstAxisTest helper that supplies axis as a constant initializer so the operator's shape-inference function is exercised (the existing tests pass axis as a runtime input, which skips that path). |
Details
- For an output of
rank + 1dimensions, a negativeaxismust be normalized as
axis + (rank + 1). The previousrank + axis - 1formula produced a negative
insertion index for the most-negative valid axes, which was then used to index
the protobuf dimension list out of bounds. - The axis value was previously read with
int32_data()[0]. When the value is
stored asraw_data(the common encoding for serialized models and the one
produced by the test harness),int32_data()is empty and the access is out of
bounds.ParseScalardecodes either encoding and validates the count.
Testing
- Built
onnxruntime_provider_testand ran the ExpandDims suite:
./onnxruntime_provider_test --gtest_filter="ExpandDimsTest.*"— all 6 tests pass. - Confirmed the new regression test fails (process aborts) without the fix and
passes with it. - Existing positive/negative out-of-range and kernel tests are unchanged.
Checklist
- Tests added/updated
- Documentation updated (if applicable)
- No breaking changes
- CI passes