1mod adam;
2mod momentum;
3mod rms_prop;
4mod sgd;
5
6use luma_io::lpk::LumaPack;
7use luma_tensor::{Device, GradStore};
8
9pub use adam::*;
10pub use momentum::*;
11pub use rms_prop::*;
12pub use sgd::*;
13
14pub trait Optimizer {
15 type Device: Device;
16
17 fn get_lr(&self) -> f64;
18 fn set_lr(&mut self, lr: f64);
19 fn step(&mut self, grads: &GradStore<Self::Device>) -> luma_tensor::Result<()>;
20 fn state_dict(&self) -> luma_tensor::Result<LumaPack<Self::Device>>;
21 fn load_state_dict(&mut self, pack: &LumaPack<Self::Device>) -> luma_tensor::Result<()>;
22}