Fix negative-axis handling in ExpandDims shape inference

Merged
TI
4mo ago
+65-2

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

FileChange
onnxruntime/core/graph/contrib_ops/contrib_defs.ccNormalize 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

FileChange
onnxruntime/test/contrib_ops/expand_dims_test.ccAdd 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 + 1 dimensions, a negative axis must be normalized as
    axis + (rank + 1). The previous rank + axis - 1 formula 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 as raw_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. ParseScalar decodes either encoding and validates the count.

Testing

  • Built onnxruntime_provider_test and 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