Fix reduction operator shape inference at opset 18 and later

[Kayvan Zahiri (Kayvan-Zahiri)]

Two shape inference defects for reduction operators at opset 18 and later, where axes
became an input rather than an attribute.

Seven operators lose their output shape entirely. _onnx_infer_single_node forwards
initializers into its temp graph only for Unsqueeze, so ReduceL1, ReduceL2,
ReduceLogSum, ReduceLogSumExp, ReduceMax, ReduceMin and ReduceSumSquare reach
onnx.shape_inference with no axes value and come back with nothing. ReduceSum,
ReduceMean and ReduceProd escape because they have dedicated handlers.

# ReduceL2 over [2, 3, 4], axes=[1] as an initializer, keepdims=0, opset 18
# inferred output shape: none at all      expected [2, 4]

ReduceSum ignores noop_with_empty_axes. An empty axes tensor is not None, so
the per-axis loop matches nothing and the input shape is copied through. The spec reduces
every axis unless noop_with_empty_axes=1. ReduceMean routes to the same handler at
opset 18.

# ReduceSum over [2, 3, 4], empty axes, keepdims=0, noop_with_empty_axes=0
# inferred [2, 3, 4]      expected [] (scalar)

noop_with_empty_axes=1 still passes the shape through unchanged, and an explicit
non-empty axes list is unaffected. Both covered in the added tests.

onnxruntime_test_python_symbolic_shape_infer.py goes from 10 failed / 31 passed to
2 failed / 31 passed with 9 subtests passing. The two remaining failures are the
pre-existing TestSymbolicShapeInferenceForSlice step tests, unrelated and failing on main.

[azure-pipelines]

Azure Pipelines:
There may be pipelines that require an authorized user to comment /azp run to run.

[Copilot]

🟡 Changes recommended

Opset-18 ReduceProd remains unhandled and can still lose its output shape.

Once you've addressed the issues Copilot identified, you can request another Copilot review.

Pull request overview

Fixes symbolic shape inference for reduction operators using axes inputs.

Changes:

• Forwards reduction axes initializers to ONNX shape inference.
• Handles empty axes according to noop_with_empty_axes.
• Adds opset-18 regression tests.

File summaries

File	Description
onnxruntime/python/tools/symbolic_shape_infer.py	Updates reduction shape inference behavior.
onnxruntime/test/python/onnxruntime_test_python_symbolic_shape_infer.py	Adds axes-input and empty-axes tests.

Review details

• Files reviewed: 2/2 changed files
• Comments generated: 1
• Review effort level: Balanced

---

💡 Configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[Kayvan Zahiri (Kayvan-Zahiri)]

Correct, and I verified it before changing anything. Fixed in the latest commit.

The reason it applies even though ReduceProd has a dispatched handler is the call order
in _infer_impl: _onnx_infer_single_node(node) runs at line 2832, then
self.dispatcher_[node.op_type](node) at 2835. So the forwarding list is consulted for
dispatched operators too. _infer_ReduceProd only populates sympy_data_ for the
keep_dims == 0 and axes == [0] case reached through the legacy axes attribute, and never
writes the output shape, so at opset 18 the shape came only from onnx shape inference,
which had no axes value.

Confirmed on the branch before the change:

ReduceProd   -> NO SHAPE     expected [2, 4]
ReduceSum    -> [2, 4]
ReduceMean   -> [2, 4]
ReduceMax    -> [2, 4]

ReduceSum and ReduceMean escape because their handlers do set the shape. That makes
eight operators needing the forwarding and two with working handlers, which matches the ten
in the spec. All ten now infer [2, 4] for a [2, 3, 4] input reduced over axis 1 with
keepdims=0.

ReduceProd is added to the parameterized test. The opset 13 attribute path is unchanged,
still inferring [] for axes=[0], keepdims=0 on a rank-1 input.

The suite goes from 11 failed / 31 passed to 2 failed / 31 passed with 10 subtests passing.
The two remaining are the pre-existing TestSymbolicShapeInferenceForSlice step tests,
which fail on main and are unrelated.