[Lanre Shittu (Shizoqua)]

Description

_infer_LayerNormalization in symbolic_shape_infer.py now takes the type of the Mean and InvStdDev outputs from the stash_type attribute (default float) instead of from input X. The same handler covers SimplifiedLayerNormalization, whose inv_std_var output follows stash_type as well.

This matches the ONNX LayerNormalization schema (U is float or bfloat16, set by stash_type), the onnxruntime LayerNormalization and SimplifiedLayerNormalization schemas, and the CPU kernel, which registers U as float.

Added test_layer_norm_stats_type_follows_stash_type to onnxruntime_test_python_symbolic_shape_infer.py. It covers double input with the default stash_type, float input with stash_type=bfloat16, and SimplifiedLayerNormalization with double input.

Motivation and Context

Fixes #33982

For a double input the tool wrote tensor(double) for these outputs, so a valid model failed to load in onnxruntime after shape inference:

Type Error: Type (tensor(double)) of output arg (mean) of node () does not match expected type (tensor(float)).

Testing, on Linux with the onnxruntime 1.31.1 CPU wheel as the runtime and the Python tooling from this branch:

• The new test fails on main (3 subtests: double where float is expected, and stash_type=bfloat16 ignored) and passes with the fix.
• python -m unittest onnxruntime_test_python_symbolic_shape_infer: 49 tests, 47 pass. The 2 failures (test_symbolic_step, test_symbolic_negative_step) also fail on main with sympy 1.14, which CI pins, and pass with sympy 1.13.3. They are unrelated to this change.
• The reproduction model from the issue loads in onnxruntime after symbolic shape inference.
• onnxruntime/test/python/quantization: 541 passed, 18 skipped.
• lintrunner on the changed files: no issues

[azure-pipelines]

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

[Copilot]

🟢 Approval recommended

No active review issues are reported for the change aligning normalization statistics type inference with stash_type.

No open issues 🎉

What changed in this PR

Fixes #33982 by aligning symbolic shape inference for normalization statistics with the schemas’ stash_type attribute, preventing incorrect output types for double-precision inputs.

Changes:

• Derives statistics output types from stash_type, defaulting to float, for both normalization operators.
• Adds regression coverage for double inputs and explicit bfloat16 statistics.

File	Description
onnxruntime/​test/​python/​onnxruntime_test_python_symbolic_shape_infer.py	Adds three cases checking inferred output types and shapes.
onnxruntime/​python/​tools/​symbolic_shape_infer.py	Uses stash_type for statistics in the shared normalization handler.

_{Review effort: Balanced}