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

pub mod attention;
pub mod block;
pub mod decoder;
pub mod encoder;
pub mod feedforward;
pub mod flash_attn;
pub mod fused_attn;
pub mod norm;
pub mod patch_embed;
pub mod pos_embed;

use burn::module::{Param, ParamId};
use burn::nn::Linear;
use burn::prelude::*;

/// Create a [`Linear`] layer with zero-filled weights.
/// Weights will be overwritten from a safetensors file.
pub fn linear_zeros<B: Backend>(
    d_input: usize,
    d_output: usize,
    bias: bool,
    device: &B::Device,
) -> Linear<B> {
    let weight = Param::initialized(
        ParamId::new(),
        Tensor::zeros([d_input, d_output], device),
    );
    let bias = if bias {
        Some(Param::initialized(
            ParamId::new(),
            Tensor::zeros([d_output], device),
        ))
    } else {
        None
    };
    Linear { weight, bias }
}