use burn::module::{Module, Param};
use burn::nn::Initializer;
use burn::tensor::backend::Backend;
use burn::tensor::{Float, Tensor};

#[derive(Module, Debug)]
pub struct RMSNorm<B: Backend> {
    pub weight: Param<Tensor<B, 1, Float>>,
    pub eps: f64,
}

impl<B: Backend> RMSNorm<B> {
    pub fn new(dim: usize, eps: f64, device: &B::Device) -> Self {
        let weight = Initializer::Ones.init([dim], device);
        Self {
            weight,
            eps,
        }
    }

    pub fn forward<const D: usize>(&self, x: Tensor<B, D, Float>) -> Tensor<B, D, Float> {
        let variance = x.clone().powf_scalar(2.0).mean_dim(D - 1);
        let rms = (variance + self.eps).sqrt();
        let x_norm = x / rms;

        let weight = self.weight.val().clone().unsqueeze();
        x_norm * weight
    }
}