optirs_core/optimizers/
mod.rs

1// Optimization algorithms for machine learning
2//
3// This module provides various optimization algorithms commonly used in machine learning,
4// such as Stochastic Gradient Descent (SGD), Adam, RMSprop, and others.
5
6use scirs2_core::ndarray::{Array, Dimension, ScalarOperand};
7use scirs2_core::numeric::{Float, ToPrimitive};
8use std::fmt::Debug;
9
10use crate::error::{OptimError, Result};
11
12/// Fallibly converts any primitive numeric value (an `f64` literal, a `usize`
13/// step count, ...) into the optimizer's generic scalar type `A`.
14///
15/// Centralizes what used to be `A::from(x).expect("unwrap failed")` call
16/// sites across the optimizer implementations: instead of panicking, a type
17/// that genuinely cannot represent `x` now produces an honest [`OptimError`].
18pub(crate) fn cast_scalar<A: Float, T: ToPrimitive>(value: T) -> Result<A> {
19    A::from(value).ok_or_else(|| {
20        OptimError::InvalidConfig(
21            "failed to convert a numeric value to the optimizer's scalar type".to_string(),
22        )
23    })
24}
25
26/// Fallibly converts the optimizer's generic scalar type `A` into `f64`.
27///
28/// Centralizes what used to be `x.to_f64().expect("unwrap failed")` call
29/// sites used for hyperparameter validation.
30pub(crate) fn scalar_to_f64<A: Float>(value: A) -> Result<f64> {
31    value.to_f64().ok_or_else(|| {
32        OptimError::InvalidConfig(
33            "failed to convert the optimizer's scalar type to f64".to_string(),
34        )
35    })
36}
37
38/// Trait that defines the interface for optimization algorithms
39pub trait Optimizer<A, D>
40where
41    A: Float + ScalarOperand + Debug,
42    D: Dimension,
43{
44    /// Updates parameters using the given gradients
45    ///
46    /// # Arguments
47    ///
48    /// * `params` - The current parameter values
49    /// * `gradients` - The gradients of the parameters
50    ///
51    /// # Returns
52    ///
53    /// The updated parameters
54    fn step(&mut self, params: &Array<A, D>, gradients: &Array<A, D>) -> Result<Array<A, D>>;
55
56    /// Gets the current learning rate
57    fn get_learning_rate(&self) -> A;
58
59    /// Sets a new learning rate
60    fn set_learning_rate(&mut self, learning_rate: A);
61
62    /// Updates multiple parameter arrays at once
63    ///
64    /// # State contract
65    ///
66    /// Position `i` in `params_list` identifies parameter tensor `i` and **must** get
67    /// its own optimizer state (moments, accumulators, velocities and any per-tensor
68    /// timestep). The caller is expected to pass the tensors in a stable order across
69    /// calls, exactly like PyTorch's parameter groups.
70    ///
71    /// The default implementation below simply forwards to [`Optimizer::step`], which
72    /// is only correct for *stateless* optimizers. Every stateful optimizer in this
73    /// crate overrides `step_list` and routes each index to a dedicated state slot
74    /// (see e.g. `Adam::step_indexed`). Implementors of new stateful optimizers must
75    /// do the same: relying on the default makes all tensors share one state slot, so
76    /// they reset each other on every shape change and their bias correction advances
77    /// once per tensor instead of once per step.
78    ///
79    /// # Arguments
80    ///
81    /// * `params_list` - List of parameter arrays
82    /// * `gradients_list` - List of gradient arrays corresponding to the parameters
83    ///
84    /// # Returns
85    ///
86    /// Updated parameter arrays
87    fn step_list(
88        &mut self,
89        params_list: &[&Array<A, D>],
90        gradients_list: &[&Array<A, D>],
91    ) -> Result<Vec<Array<A, D>>> {
92        if params_list.len() != gradients_list.len() {
93            return Err(OptimError::InvalidConfig(format!(
94                "Number of parameter arrays ({}) does not match number of gradient arrays ({})",
95                params_list.len(),
96                gradients_list.len()
97            )));
98        }
99
100        let mut results = Vec::with_capacity(params_list.len());
101        for (params, grads) in params_list.iter().zip(gradients_list.iter()) {
102            results.push(self.step(params, grads)?);
103        }
104        Ok(results)
105    }
106}
107
108// Import specific optimizers
109mod adabound;
110mod adadelta;
111mod adagrad;
112mod adam;
113mod adamw;
114mod grouped_adam;
115mod lamb;
116mod lars;
117mod lbfgs;
118mod lion;
119mod lookahead;
120mod maml;
121mod meta_sgd;
122mod ntm_optimizer;
123mod radam;
124mod ranger;
125mod reptile;
126mod rmsprop;
127mod sam;
128mod sgd;
129mod sgd_simd;
130mod sparse_adam;
131
132// Re-export specific optimizers
133pub use adabound::AdaBound;
134pub use adadelta::AdaDelta;
135pub use adagrad::Adagrad;
136pub use adam::Adam;
137pub use adamw::AdamW;
138pub use grouped_adam::GroupedAdam;
139pub use lamb::LAMB;
140pub use lars::LARS;
141pub use lbfgs::LBFGS;
142pub use lion::Lion;
143pub use lookahead::Lookahead;
144pub use maml::{MAMLVariant, TaskBatch, MAML};
145pub use meta_sgd::MetaSGD;
146pub use ntm_optimizer::{AddressingMode, NtmConfig, NtmOptimizer};
147pub use radam::RAdam;
148pub use ranger::Ranger;
149pub use reptile::ReptileOptimizer;
150pub use rmsprop::RMSprop;
151pub use sam::SAM;
152pub use sgd::SGD;
153pub use sgd_simd::SimdSGD;
154pub use sparse_adam::{SparseAdam, SparseGradient};
optirs_core/optimizers/mod.rs

optirs_core/optimizers/
mod.rs
