Do not fuse LayerNorm when ReduceMean has keepdims=0

[guptaishaan]

Fixes #30513

LayerNormFusion and SimplifiedLayerNormFusion match the ReduceMean nodes without checking keepdims. When keepdims=0 the reduced dim is dropped, so the following Sub and Div broadcast along a different axis and the sub-graph is not a layer norm, but the fusion replaces it anyway.

Two ways that goes wrong today:

• Invalid graph. In the reported model the Div output is rank 1, so the scale-selection loop picks it as the scale input of the node that replaces the Div itself. Session creation then fails with Node input 'nm' is not a graph input, initializer, or output of a previous node.
• Silently wrong results. With shapes where broadcasting happens to line up, the fusion succeeds and the model computes something different from the original graph.

The fix checks keepdims at the three ReduceMean match sites and skips the fusion when it is 0. Layer norm exports leave keepdims at its default of 1, so nothing that fused before stops fusing.

Verified on Linux x64, CPU EP, built from this branch with gcc 13:

• The model from the issue loads instead of raising InvalidArgument.
• Two reduced models that previously fused under keepdims=0 now produce output matching the same model run with graph optimizations disabled. Before the change they did not match.
• New tests LayerNormFusion_ReduceMeanNoKeepDims_NoFusion and SimplifiedLayerNormFusion_ReduceMeanNoKeepDims_NoFusion fail before the change and pass after.
• onnxruntime_test_all --gtest_filter='GraphTransformationTests.*': 362 passed, 15 skipped (WebGPU), 0 failed. The full onnxruntime_test_all: 1884 passed, 17 skipped, 0 failed. clang-format 20.1.8 clean on both changed files.

Not verified: non-CPU execution providers, Windows and macOS builds, and training-enabled builds (ENABLE_TRAINING_CORE).

Thanks to @XYZboom for the report and the reproducer.

[azure-pipelines]

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

[Copilot]

Pull request overview

This PR fixes an optimizer correctness bug in LayerNormFusion and SimplifiedLayerNormFusion where ReduceMean nodes with keepdims=0 could be incorrectly matched and fused into a layer-norm op, even though dropping the reduced dimension changes broadcasting semantics and can make the subgraph non-equivalent (or even produce an invalid graph).

Changes:

• Add a keepdims guard for ReduceMean at all LayerNorm fusion match sites so fusion is skipped when keepdims=0.
• Add new graph transformer tests to ensure no fusion occurs for keepdims=0 patterns in both the full and simplified LayerNorm fusions.

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated no comments.

File	Description
onnxruntime/core/optimizer/layer_norm_fusion.cc	Adds a helper to require ReduceMean keeps reduced dims, and applies it at each fusion match site to prevent incorrect fusions.
onnxruntime/test/optimizer/graph_transform_test_layernorm.cc	Adds regression tests verifying that keepdims=0 ReduceMean patterns do not fuse into LayerNormalization / SimplifiedLayerNormalization.

[guptaishaan]

@microsoft-github-policy-service agree

[guptaishaan]

Thanks for the review. The automated summary matches the change, and it did not raise anything, so no code change here.

To restate the scope for a human reviewer: the guard is only on keepdims, applied at the three ReduceMean match sites in LayerNormFusion and SimplifiedLayerNormFusion. keepdims defaults to 1 and layer norm exports leave it there, so nothing that fused before stops fusing; only the keepdims=0 sub-graphs, which are not layer norms, are now skipped.

Happy to adjust anything. This is my first PR here, so let me know if the Azure Pipelines runs need authorizing.