Function burn_rms_norm_nd 

pub fn burn_rms_norm_nd<const D: usize>(
    x: Tensor<D>,
    w: Tensor<1>,
    eps: f32,
) -> Tensor<D>

RMSNorm over the last dim for rank-D input.