LayerNormalization broadcast (limited support for axis=2)

[tianleiwu]

Description

Spec of LayerNormalization supports broadcasting (tensors Scale and B should be unidirectional broadcastable to tensor X).
https://onnx.ai/onnx/operators/onnx__LayerNormalization.html
However, current implementation only allow scale and bias size to be X.shape()[axis:].

Example of input tensors that normalized with axis=2:

X shape	Scale shape	B shape	Before	After
(B, S, D)	(D)	(D)	Supported	Supported
(B, S, D)	(1, 1, D)	(1, 1, D)	Supported	Supported
(B, S, D)	(B, 1, D)	(B, 1, D)	Not Supported	Supported
(B, S, D)	(1, S, D)	(1, S, D)	Not Supported	Supported
(B, S, D)	(B, S, D)	(B, S, D)	Not Supported	Supported

Here we add limited support: axis=2; scale/bias has same shape; scale/bias/X have same number of dimensions. It could support common use case in LLM and vision models.

Motivation and Context

Support Stable Diffusion 3.x and Flux model.