fuse_layernorm_or_instancenorm: check gamma / beta shapes, not just their rank

[LeSingh1]

_try_apply_transform decides a pattern is a layer_norm when gamma and beta have rank len(axes), and an instance_norm when they squeeze to rank 1:

if gamma_rank == len(axes) and beta_rank == len(axes):
    ...
if len(np.squeeze(gamma_var.val).shape) == 1 and len(np.squeeze(beta_var.val).shape) == 1:
    is_instancenorm = True

Neither looks at the shape, so a gamma that merely broadcasts over those axes is fused into an op it does not fit. Both cases break the conversion of a program that is valid before the pass:

layer_norm. axes=[-1] on a (1, 2, 5) input with gamma of shape (1,) — a normalization with a scalar affine — has rank 1, so it is fused, and layer_norm.type_inference (which requires gamma.shape == x.shape[axes]) immediately raises:

ValueError: Expect shape [5] for gamma, but get shape (1,) instead

instance_norm. axes=[-2, -1] on a (1, 3, 4, 5) input with gamma of shape (1, 1, 4, 1) squeezes to (4,), so it is fused into instance_norm(gamma=<len 4>) — moving a factor that applied to H onto the channel axis, which is 3 long. instance_norm.type_inference does not check the length, so this reaches the backend:

RuntimeError: Error compiling model: ... Dimension 0 of tensor parameter beta[0]
has unexpected length 4; expected 3.

Fix

Require gamma and beta to have shape x.shape[axes] for layer_norm (reusing layer_norm._is_compatible_shape, so symbolic dims stay allowed), and to be C long for instance_norm, where C is x.shape[-3] for axes=[-2, -1] and x.shape[-1] for the channel-last axes=[-3, -2] case.

Patterns that do fit are fused exactly as before; the pass just fuses fewer of the ones it cannot represent.

Testing

Three tests in TestFuseLayerNormOrInstanceNorm: one per case above, plus a control that a proper per-channel gamma is still fused into instance_norm. The two negative tests fail on main. The existing tests in that class are unchanged and still pass.