29 lines
811 B
Rust
29 lines
811 B
Rust
|
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
|
||
|
|
}
|
||
|
|
}
|