ruda-nn 0.21.23

use ruda_model::tensor::{DType, FloatDType};


use ruda_model::config::Config;
use ruda_model::module::Initializer;
use ruda_model::module::Module;
use ruda_model::module::Param;
use ruda_model::module::{Content, DisplaySettings, ModuleDisplay};
use ruda_model::tensor::Tensor;
use ruda_model::tensor::backend::Backend;

/// Configuration to create a [RMS Norm](RmsNorm) layer using the [init function](RmsNormConfig::init).
#[derive(Config, Debug)]
pub struct RmsNormConfig {
    /// The size of the input features.
    pub d_model: usize,
    /// A value required for numerical stability. Default: 1e-5
    #[config(default = 1e-5)]
    pub epsilon: f64,
}

impl RmsNormConfig {
    /// Initialize a new [RMS Norm](RmsNorm) module.
    ///
    /// # Panics
    ///
    /// Panics if `epsilon` is not positive.
    pub fn init<B: Backend>(&self, device: &B::Device) -> RmsNorm<B> {
        assert!(self.epsilon > 0.0, "epsilon must be positive.");

        let gamma = Initializer::Ones.init([self.d_model], device);

        RmsNorm {
            gamma,
            epsilon: self.epsilon,
        }
    }
}

/// Applies RMS Normalization over an input tensor along the last dimension.
///
/// `Y = X / sqrt(mean(X^2) + eps) * gamma`
///
/// Where:
/// - `X` is the input tensor
/// - `Y` is the output tensor
/// - `gamma` is the learnable weight
/// - `mean` is the mean operation
/// - `eps` is a small value to avoid division by zero.
///
/// Should be created using the [RmsNormConfig](RmsNormConfig) configuration.
#[derive(Module, Debug)]
#[module(custom_display)]
pub struct RmsNorm<B: Backend> {
    /// The learnable parameter to scale the normalized tensor
    pub gamma: Param<Tensor<B, 1>>,
    /// A value required for numerical stability
    pub epsilon: f64,
}

impl<B: Backend> RmsNorm<B> {
    pub fn forward_with_compute_dtype<const D: usize>(
        &self,
        input: Tensor<B, D>,
        dtype: FloatDType,
    ) -> Tensor<B, D> {
        let output_dtype = input.dtype();
        let dtype: DType = dtype.into();
        let input = input.cast(dtype);
        let rms = (input.clone().square().mean_dim(D - 1) + self.epsilon).sqrt();
        ((input / rms) * self.gamma.val().cast(dtype).unsqueeze()).cast(output_dtype)
    }
    /// Applies the forward pass on the input tensor.
    ///
    /// See the [RmsNorm](RmsNorm) documentation for more information.
    ///
    /// # Shapes
    ///
    /// - input: `[..., any, d_model]`
    /// - output: `[..., any, d_model]`
    pub fn forward<const D: usize>(&self, x: Tensor<B, D>) -> Tensor<B, D> {
        // Calculate the root-mean-square norm of the input tensor along the last dimension
        let dtype = x.dtype();
        let rms = (x.clone().cast(DType::F32).square().mean_dim(D - 1) + self.epsilon).sqrt();
        (x / rms.cast(dtype)) * self.gamma.val().unsqueeze()
    }
}

impl<B: Backend> ModuleDisplay for RmsNorm<B> {
    fn custom_settings(&self) -> Option<DisplaySettings> {
        DisplaySettings::new()
            .with_new_line_after_attribute(false)
            .optional()
    }

    fn custom_content(&self, content: Content) -> Option<Content> {
        let [d_model] = self.gamma.shape().dims();
        content
            .add("d_model", &d_model)
            .add("epsilon", &self.epsilon)
            .optional()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::TestBackend;
    use alloc::format;
    use ruda_model::tensor::TensorData;
    use ruda_model::tensor::{Tolerance, ops::FloatElem};
    type FT = FloatElem<TestBackend>;

    #[test]
    fn rms_norm_forward() {
        let device = Default::default();
        let module = RmsNormConfig::new(3)
            .with_epsilon(1e-5)
            .init::<TestBackend>(&device);

        let input = Tensor::arange(0..9, &device).float().reshape([3, 3]);
        let output = module.forward(input);

        let expected = TensorData::from([
            [0.0000, 0.7746, 1.5492],
            [0.7348, 0.9798, 1.2247],
            [0.8514, 0.9933, 1.1352],
        ]);
        output
            .to_data()
            .assert_approx_eq::<FT>(&expected, Tolerance::default());
    }

    #[cfg(feature = "std")]
    #[test]
    fn rms_norm_compute_dtype_preserves_storage_and_original_gradients() {
        use crate::TestAutodiffBackend as B;
        let device = Default::default();
        for dtype in [DType::F16, DType::BF16] {
            let module = RmsNormConfig::new(3).init::<B>(&device);
            let input = Tensor::<B, 2>::from_floats([[1., 2., 4.], [3., -1., 2.]], &device)
                .cast(dtype).require_grad();
            let reference = module.forward(input.clone().cast(DType::F32)).cast(dtype);
            let output = module.forward_with_compute_dtype(input.clone(), FloatDType::F32);
            assert_eq!(output.dtype(), dtype);
            output.clone().cast(DType::F32).to_data().assert_approx_eq::<f32>(
                &reference.clone().cast(DType::F32).to_data(), Tolerance::absolute(1e-6));
            let expected_grads = reference.square().sum().backward();
            let grads = output.square().sum().backward();
            input.grad(&grads).unwrap().cast(DType::F32).to_data().assert_approx_eq::<f32>(
                &input.grad(&expected_grads).unwrap().cast(DType::F32).to_data(), Tolerance::absolute(1e-6));
            module.gamma.val().grad(&grads).unwrap().to_data().assert_approx_eq::<f32>(
                &module.gamma.val().grad(&expected_grads).unwrap().to_data(), Tolerance::absolute(1e-6));
        }
    }

    #[test]
    fn display() {
        let config = RmsNormConfig::new(6);
        let layer_norm = config.init::<TestBackend>(&Default::default());

        assert_eq!(
            format!("{layer_norm}"),
            "RmsNorm {d_model: 6, epsilon: 0.00001, params: 6}"
        );
    }
}