Adding RMSNorm support to arbitrary x and normalized_dim shapes (#9966)

Summary:

In D72014553, we were adding initial support for RMS norm for an input in 3 or 4 dimensions, and a weight of dimension 1 (same size as x[:-1])

In this diff, we allow for:
- input of arbitrary shape
- shape broadcasting of w (w must have dim <= 1)

Differential Revision: D72484196

Author: ThomasJannaud

backends/cadence/aot/ops_registrations.py

@@ -139,7 +139,9 @@
     "int in_zero_point, bool channel_last=False) -> (Tensor out)"
 )
 lib.define("linalg_vector_norm(Tensor X) -> (Tensor Y)")
-lib.define("rms_norm(Tensor X, float eps, Tensor W) -> (Tensor Y)")
+lib.define(
+    "rms_norm(Tensor input, SymInt[] normalized_shape, Tensor? weight=None, float? eps=None) -> (Tensor Y)"
+)
 lib.define(
     "transposed_im2row(Tensor input, int[2] kernel_size, int[2] dilation, int[2] padding, int[2] stride, "
     "int[2] output_padding, Tensor in_zero_point, bool channel_last=False) -> (Tensor out)"
@@ -212,7 +214,7 @@
 )
 lib.define("linalg_vector_norm.out(Tensor X, *, Tensor(a!) out) -> Tensor(a!)")
 lib.define(
-    "rms_norm.out(Tensor X, float eps, Tensor W, *, Tensor(a!) out) -> Tensor(a!)"
+    "rms_norm.out(Tensor input, SymInt[] normalized_shape, Tensor? weight=None, float? eps=None, *, Tensor(a!) out) -> Tensor(a!)"
 )
 lib.define(
     "quantized_fully_connected.out(Tensor src, Tensor weight, Tensor bias, int src_zero_point, "