Fix LayerNormalization stats dtype in symbolic shape inference

onnxruntime/python/tools/symbolic_shape_infer.py

@@ -2560,9 +2560,7 @@ def _infer_LayerNormalization(self, node):  # noqa: N802
                 rank = len(x_shape)
                 axis = handle_negative_axis(axis, rank)
                 mean_shape = x_shape[:axis] + [1 for _ in range(rank - axis)]
-                mean_dtype = self.known_vi_[node.input[0]].type.tensor_type.elem_type
-                if mean_dtype == onnx.TensorProto.FLOAT16 or mean_dtype == onnx.TensorProto.BFLOAT16:
-                    mean_dtype = onnx.TensorProto.FLOAT
+                mean_dtype = get_attribute(node, "stash_type", onnx.TensorProto.FLOAT)
                 vi = self.known_vi_[node.output[1]]
                 vi.CopyFrom(helper.make_tensor_value_info(node.output[1], mean_dtype, mean_shape))
                 if len(node.output) > 2:

onnxruntime/test/python/onnxruntime_test_python_symbolic_shape_infer.py

@@ -1484,6 +1484,37 @@ def test_matmulnbits(self):
         ]
         self._check_shapes(graph, inferred.graph, expected_shapes)
 
+    def test_layer_norm_stats_type_follows_stash_type(self):
+        cases = [
+            ("LayerNormalization", TensorProto.DOUBLE, None, TensorProto.FLOAT),
+            ("LayerNormalization", TensorProto.FLOAT, TensorProto.BFLOAT16, TensorProto.BFLOAT16),
+            ("SimplifiedLayerNormalization", TensorProto.DOUBLE, None, TensorProto.FLOAT),
+        ]
+        for op_type, x_type, stash_type, expected_type in cases:
+            with self.subTest(op_type=op_type, x_type=x_type, stash_type=stash_type):
+                inputs = ["X", "scale"] if op_type == "SimplifiedLayerNormalization" else ["X", "scale", "bias"]
+                outputs = (
+                    ["Y", "inv_std_var"] if op_type == "SimplifiedLayerNormalization" else ["Y", "mean", "inv_std"]
+                )
+                attrs = {} if stash_type is None else {"stash_type": stash_type}
+                node = helper.make_node(op_type, inputs, outputs, axis=-1, **attrs)
+                initializers = [helper.make_tensor(name, x_type, [4], [1.0] * 4) for name in inputs[1:]]
+                graph = helper.make_graph(
+                    [node],
+                    "LayerNorm_Test",
+                    [helper.make_tensor_value_info("X", x_type, ["batch", 4])],
+                    [helper.make_tensor_value_info(name, TensorProto.UNDEFINED, None) for name in outputs],
+                    initializers,
+                )
+                model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 17)])
+
+                inferred = SymbolicShapeInference.infer_shapes(model, auto_merge=True)
+
+                expected_shapes = [helper.make_tensor_value_info("Y", x_type, ["batch", 4])] + [
+                    helper.make_tensor_value_info(name, expected_type, ["batch", 1]) for name in outputs[1:]
+                ]
+                self._check_shapes(graph, inferred.graph, expected_shapes)
+
     def test_qlinear_binary(self):
         """
         Test ONNX QLinearAdd op ('com.microsoft' domain). .