eegdino 0.1.0

EEG-DINO 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
23
24
25
26
27
28
29
30
31
32
33
34

/// Feed-forward MLP with GELU activation.
///
/// Matches the Python `Mlp` class in `models/transformer.py`:
///   fc1(x) → GELU → fc2(x)
///
/// Dropout is omitted for inference (eval mode).
use burn::prelude::*;
use burn::nn::Linear;

use super::linear_zeros;

#[derive(Module, Debug)]
pub struct Mlp<B: Backend> {
    pub fc1: Linear<B>,
    pub fc2: Linear<B>,
}

impl<B: Backend> Mlp<B> {
    pub fn new(in_features: usize, hidden_features: usize, device: &B::Device) -> Self {
        Self {
            fc1: linear_zeros::<B>(in_features, hidden_features, true, device),
            fc2: linear_zeros::<B>(hidden_features, in_features, true, device),
        }
    }

    /// Forward pass.
    ///
    /// x: [B, N, in_features] → [B, N, in_features]
    pub fn forward(&self, x: Tensor<B, 3>) -> Tensor<B, 3> {
        let x = self.fc1.forward(x);
        let x = burn::tensor::activation::gelu(x);
        self.fc2.forward(x)
    }
}