Function rms_norm 

pub fn rms_norm<B: Backend, const D: usize>(
    x: Tensor<B, D>,
    weight: Tensor<B, 1>,
    eps: f64,
) -> Tensor<B, D>

y = x / rms(x) * w where rms is taken over the last dimension and eps is added inside the square root.

The reduction runs in f32 for f16 stability (no-op in f32 builds).