lattice-tune 0.2.2

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

//! GPU buffer state management

/// GPU buffer for optimizer state
pub struct OptimizerState {
    /// First moment (Adam)
    pub m: wgpu::Buffer,
    /// Second moment (Adam)
    pub v: wgpu::Buffer,
    /// Velocity (SGD momentum)
    pub velocity: wgpu::Buffer,
    /// Current timestep
    pub t: u32,
}

/// Layer gradients stored on GPU
pub struct LayerGradients {
    pub weight_grads: wgpu::Buffer,
    pub bias_grads: wgpu::Buffer,
    pub optimizer_state: OptimizerState,
    pub num_weights: usize,
    pub num_biases: usize,
}