lattice-tune 0.3.0

Training infrastructure for Lattice neural models
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
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60

//! GPU shader uniform structures

/// Uniforms for Adam optimizer shader
#[repr(C)]
#[derive(Copy, Clone, Debug, bytemuck::Pod, bytemuck::Zeroable)]
pub struct AdamUniforms {
    pub size: u32,
    pub learning_rate: f32,
    pub beta1: f32,
    pub beta2: f32,
    pub epsilon: f32,
    pub t: f32,
    pub _pad0: u32,
    pub _pad1: u32,
}

/// Uniforms for AdamW optimizer shader
#[repr(C)]
#[derive(Copy, Clone, Debug, bytemuck::Pod, bytemuck::Zeroable)]
pub struct AdamWUniforms {
    pub size: u32,
    pub learning_rate: f32,
    pub beta1: f32,
    pub beta2: f32,
    pub epsilon: f32,
    pub weight_decay: f32,
    pub t: f32,
    pub _pad: u32,
}

/// Uniforms for SGD momentum shader
#[repr(C)]
#[derive(Copy, Clone, Debug, bytemuck::Pod, bytemuck::Zeroable)]
pub struct SgdMomentumUniforms {
    pub size: u32,
    pub learning_rate: f32,
    pub momentum: f32,
    pub _pad: u32,
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_adam_uniforms_size() {
        // Verify struct is properly aligned for GPU
        assert_eq!(std::mem::size_of::<AdamUniforms>(), 32);
    }

    #[test]
    fn test_adamw_uniforms_size() {
        assert_eq!(std::mem::size_of::<AdamWUniforms>(), 32);
    }

    #[test]
    fn test_sgd_momentum_uniforms_size() {
        assert_eq!(std::mem::size_of::<SgdMomentumUniforms>(), 16);
    }
}