brainharmony 0.1.0

Brain-Harmony multimodal brain foundation model — inference in Rust with Burn ML
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22

/// LayerNorm wrapper (burn 0.20.1)
///
/// Brain-Harmony uses pre-norm LayerNorm (eps=1e-6), matching PyTorch nn.LayerNorm.
use burn::nn::{LayerNorm, LayerNormConfig};
use burn::prelude::*;

#[derive(Module, Debug)]
pub struct LNorm<B: Backend> {
    pub inner: LayerNorm<B>,
}

impl<B: Backend> LNorm<B> {
    pub fn new(dim: usize, eps: f64, device: &B::Device) -> Self {
        Self {
            inner: LayerNormConfig::new(dim).with_epsilon(eps).init(device),
        }
    }

    pub fn forward<const D: usize>(&self, x: Tensor<B, D>) -> Tensor<B, D> {
        self.inner.forward(x)
    }
}