torsh_optim/
lib.rs

1//! Optimization algorithms for ToRSh
2//!
3//! This crate provides PyTorch-compatible optimizers built on top of scirs2-optim.
4//!
5//! # Features
6//!
7//! - **80+ optimizers**: Comprehensive collection including Adam, SGD, RAdam, Ranger, Lion, Sophia, and more
8//! - **Modern optimizers**: Latest research including Schedule-Free AdamW and Prodigy
9//! - **Second-order methods**: L-BFGS, Newton-CG, Trust Region, K-FAC, AdaHessian
10//! - **Learning rate schedulers**: Step, exponential, cosine annealing, one-cycle, and more
11//! - **Mixed precision training**: Full fp16/fp32 support with loss scaling
12//! - **Distributed optimization**: AsyncSGD, Elastic Averaging, Federated Learning
13//! - **Advanced features**: Gradient accumulation, fused kernels, memory-efficient implementations
14//! - **Research features**: Quantum-inspired, neuromorphic, continual learning, green AI optimizers
15//!
16//! # Quick Start
17//!
18//! ```rust,no_run
19//! use torsh_optim::prelude::*;
20//! use torsh_tensor::Tensor;
21//! use std::sync::Arc;
22//! use parking_lot::RwLock;
23//!
24//! # fn main() -> Result<(), Box<dyn std::error::Error>> {
25//! // Create parameters
26//! let params = vec![Arc::new(RwLock::new(Tensor::scalar(1.0)?))];
27//!
28//! // Create optimizer
29//! let mut optimizer = Adam::new(params, Some(0.001), None, None, None, false);
30//!
31//! // Training loop
32//! for _ in 0..100 {
33//!     // ... compute gradients ...
34//!     optimizer.step()?;
35//!     optimizer.zero_grad();
36//! }
37//! # Ok(())
38//! # }
39//! ```
40
41#![cfg_attr(not(feature = "std"), no_std)]
42// Note: These allows are necessary for maintaining compatibility with diverse optimizer implementations
43// and reducing noise from legitimate design patterns used across the codebase
44#![allow(dead_code)] // Many optimizers have internal methods not called externally
45#![allow(unused_imports)] // Conditional compilation features may leave some imports unused
46#![allow(unused_variables)] // Some optimizer variants have parameters used only in specific configurations
47#![allow(unused_mut)] // Mutability annotations required for consistency even when not always modified
48
49#[cfg(not(feature = "std"))]
50extern crate alloc;
51
52pub mod adabelief;
53pub mod adabound;
54pub mod adadelta;
55pub mod adagrad;
56pub mod adahessian;
57pub mod adam;
58pub mod adamax;
59pub mod advanced;
60pub mod asgd;
61pub mod bayesian_optimization;
62pub mod benchmarks;
63pub mod checkpointing;
64pub mod composition;
65pub mod continual_learning;
66pub mod cross_framework_validation;
67pub mod debugging;
68pub mod differential_privacy;
69pub mod distributed;
70pub mod evolutionary_strategies;
71pub mod ftrl;
72pub mod fused_kernels;
73pub mod grad_accumulation;
74pub mod gradient_free;
75pub mod green_ai;
76pub mod hyperparameter_tuning;
77pub mod kfac;
78pub mod lamb;
79pub mod lazy_updates;
80pub mod lbfgs;
81pub mod lion;
82pub mod lookahead;
83pub mod low_precision;
84pub mod lr_scheduler;
85pub mod lr_scheduler_additional;
86pub mod lr_scheduler_enhanced;
87pub mod memory_efficient;
88pub mod memory_mapped;
89pub mod mixed_precision;
90pub mod nadam;
91pub mod natural_gradient;
92pub mod neural_optimizer;
93pub mod neuromorphic;
94pub mod newton_cg;
95pub mod numerical_stability_tests;
96pub mod online_learning;
97pub mod optimizer;
98pub mod prodigy;
99pub mod quantum_inspired;
100pub mod radam;
101pub mod ranger;
102pub mod rmsprop;
103pub mod robustness;
104pub mod rprop;
105pub mod schedule_free;
106pub mod sgd;
107pub mod shampoo;
108pub mod sophia;
109pub mod sparse_adam;
110pub mod sparse_updates;
111pub mod state_dict_ops;
112pub mod stress_tests;
113pub mod trust_region;
114pub mod yellowfin;
115
116use parking_lot::RwLock;
117use std::collections::HashMap;
118use std::sync::Arc;
119use torsh_core::error::{Result, TorshError};
120use torsh_tensor::Tensor;
121
122/// Optimizer-specific error type
123#[derive(Debug, thiserror::Error)]
124pub enum OptimizerError {
125    #[error("Tensor operation failed: {0}")]
126    TensorError(#[from] torsh_core::error::TorshError),
127
128    #[error("Invalid parameter: {0}")]
129    InvalidParameter(String),
130
131    #[error("Serialization error: {0}")]
132    SerializationError(String),
133
134    #[error("IO error: {0}")]
135    IoError(#[from] std::io::Error),
136
137    #[error("Checkpoint error: {0}")]
138    CheckpointError(String),
139
140    #[error("Configuration error: {0}")]
141    ConfigError(String),
142
143    #[error("State error: {0}")]
144    StateError(String),
145
146    #[error("Invalid input: {0}")]
147    InvalidInput(String),
148
149    #[error("Numerical error: {0}")]
150    NumericalError(String),
151
152    #[error("Memory map error: {0}")]
153    MemoryMapError(String),
154}
155
156impl From<OptimizerError> for torsh_core::error::TorshError {
157    fn from(err: OptimizerError) -> Self {
158        match err {
159            OptimizerError::TensorError(e) => e,
160            OptimizerError::InvalidParameter(msg) => {
161                torsh_core::error::TorshError::InvalidArgument(msg)
162            }
163            OptimizerError::SerializationError(msg) => {
164                torsh_core::error::TorshError::SerializationError(msg)
165            }
166            OptimizerError::IoError(e) => torsh_core::error::TorshError::IoError(e.to_string()),
167            OptimizerError::CheckpointError(msg) => {
168                torsh_core::error::TorshError::RuntimeError(msg)
169            }
170            OptimizerError::ConfigError(msg) => torsh_core::error::TorshError::ConfigError(msg),
171            OptimizerError::StateError(msg) => torsh_core::error::TorshError::RuntimeError(msg),
172            OptimizerError::InvalidInput(msg) => {
173                torsh_core::error::TorshError::InvalidArgument(msg)
174            }
175            OptimizerError::NumericalError(msg) => torsh_core::error::TorshError::RuntimeError(msg),
176            OptimizerError::MemoryMapError(msg) => torsh_core::error::TorshError::RuntimeError(msg),
177        }
178    }
179}
180
181/// Result type for optimizer operations
182pub type OptimizerResult<T> = std::result::Result<T, OptimizerError>;
183
184// Version information
185pub const VERSION: &str = env!("CARGO_PKG_VERSION");
186pub const VERSION_MAJOR: u32 = 0;
187pub const VERSION_MINOR: u32 = 1;
188pub const VERSION_PATCH: u32 = 0;
189
190// Re-export scirs2 optimizer functionality
191// use scirs2::optim as sci_optim;
192
193/// Base optimizer trait
194pub trait Optimizer {
195    /// Perform a single optimization step
196    fn step(&mut self) -> OptimizerResult<()>;
197
198    /// Zero all gradients
199    fn zero_grad(&mut self);
200
201    /// Get the current learning rate
202    fn get_lr(&self) -> Vec<f32>;
203
204    /// Set the learning rate
205    fn set_lr(&mut self, lr: f32);
206
207    /// Add a parameter group
208    fn add_param_group(&mut self, params: Vec<Arc<RwLock<Tensor>>>, options: HashMap<String, f32>);
209
210    /// Get the parameter tensors managed by this optimizer.
211    ///
212    /// Returns clones of the `Arc<RwLock<Tensor>>` handles. Because they are
213    /// reference-counted, the returned handles point at the *same* underlying
214    /// tensors the optimizer updates, so callers can both read parameters and
215    /// access their gradients via [`Tensor::grad`].
216    ///
217    /// Meta-optimizers such as [`crate::lookahead::Lookahead`] and the gradient
218    /// accumulation wrappers in [`crate::grad_accumulation`] rely on this to
219    /// inspect and update the wrapped optimizer's parameters.
220    ///
221    /// # Default Implementation
222    ///
223    /// The default returns an empty vector. Concrete optimizers that own
224    /// parameter groups override this to expose their parameters, and wrapper
225    /// optimizers delegate to the optimizer they wrap. An empty result therefore
226    /// means "this optimizer does not expose parameters", *not* "this optimizer
227    /// has no parameters"; callers that require parameters must treat an empty
228    /// result as an error rather than silently doing nothing.
229    fn parameters(&self) -> Vec<Arc<RwLock<Tensor>>> {
230        Vec::new()
231    }
232
233    /// Get state dict for serialization
234    fn state_dict(&self) -> OptimizerResult<OptimizerState>;
235
236    /// Load state dict
237    fn load_state_dict(&mut self, state: OptimizerState) -> OptimizerResult<()>;
238}
239
240/// Optimizer state for serialization
241#[derive(Debug, Clone)]
242pub struct OptimizerState {
243    /// Optimizer type identifier
244    pub optimizer_type: String,
245    /// Version of the state format
246    pub version: String,
247    /// Parameter group states
248    pub param_groups: Vec<ParamGroupState>,
249    /// Per-parameter optimizer state (keyed by parameter ID)
250    pub state: HashMap<String, HashMap<String, Tensor>>,
251    /// Global optimizer state
252    pub global_state: HashMap<String, f32>,
253}
254
255/// Parameter group state
256#[derive(Debug, Clone)]
257pub struct ParamGroupState {
258    /// Learning rate for this group
259    pub lr: f32,
260    /// Additional options for this group
261    pub options: HashMap<String, f32>,
262    /// Number of parameters in this group (for validation)
263    pub param_count: usize,
264}
265
266impl OptimizerState {
267    /// Create a new empty optimizer state
268    pub fn new(optimizer_type: String) -> Self {
269        Self {
270            optimizer_type,
271            version: VERSION.to_string(),
272            param_groups: Vec::new(),
273            state: HashMap::new(),
274            global_state: HashMap::new(),
275        }
276    }
277
278    /// Validate the state structure
279    pub fn validate(&self) -> Result<()> {
280        if self.optimizer_type.is_empty() {
281            return Err(TorshError::InvalidArgument(
282                "Optimizer type cannot be empty".to_string(),
283            ));
284        }
285
286        // Check that all parameter groups are valid
287        for (i, group) in self.param_groups.iter().enumerate() {
288            if !group.lr.is_finite() || group.lr <= 0.0 {
289                return Err(TorshError::InvalidArgument(format!(
290                    "Invalid learning rate in group {}",
291                    i
292                )));
293            }
294        }
295
296        // Check that all state values are finite
297        for (param_id, param_state) in &self.state {
298            for (state_name, tensor) in param_state {
299                // For now, just check that the keys are valid
300                if param_id.is_empty() || state_name.is_empty() {
301                    return Err(TorshError::InvalidArgument(
302                        "State keys cannot be empty".to_string(),
303                    ));
304                }
305            }
306        }
307
308        Ok(())
309    }
310
311    /// Get the total number of parameters across all groups
312    pub fn total_param_count(&self) -> usize {
313        self.param_groups.iter().map(|g| g.param_count).sum()
314    }
315
316    /// Check if state is compatible with another state (same structure)
317    pub fn is_compatible_with(&self, other: &OptimizerState) -> bool {
318        self.optimizer_type == other.optimizer_type
319            && self.param_groups.len() == other.param_groups.len()
320            && self
321                .param_groups
322                .iter()
323                .zip(other.param_groups.iter())
324                .all(|(a, b)| a.param_count == b.param_count)
325    }
326}
327
328impl ParamGroupState {
329    /// Create a new parameter group state
330    pub fn new(lr: f32, param_count: usize) -> Self {
331        Self {
332            lr,
333            options: HashMap::new(),
334            param_count,
335        }
336    }
337
338    /// Create from a ParamGroup
339    pub fn from_param_group(group: &ParamGroup) -> Self {
340        Self {
341            lr: group.lr,
342            options: group.options.clone(),
343            param_count: group.params.len(),
344        }
345    }
346
347    /// Get an option value with a default
348    pub fn get_option(&self, key: &str, default: f32) -> f32 {
349        self.options.get(key).copied().unwrap_or(default)
350    }
351
352    /// Set an option value
353    pub fn set_option(&mut self, key: String, value: f32) {
354        self.options.insert(key, value);
355    }
356}
357
358/// Parameter group
359#[derive(Debug, Clone)]
360pub struct ParamGroup {
361    pub params: Vec<Arc<RwLock<Tensor>>>,
362    pub lr: f32,
363    pub options: HashMap<String, f32>,
364}
365
366/// Builder for creating parameter groups with various options
367#[derive(Debug)]
368pub struct ParamGroupBuilder {
369    params: Vec<Arc<RwLock<Tensor>>>,
370    lr: f32,
371    options: HashMap<String, f32>,
372}
373
374impl ParamGroupBuilder {
375    /// Create a new parameter group builder
376    pub fn new(lr: f32) -> Self {
377        Self {
378            params: Vec::new(),
379            lr,
380            options: HashMap::new(),
381        }
382    }
383
384    /// Add parameters to the group
385    pub fn params(mut self, params: Vec<Arc<RwLock<Tensor>>>) -> Self {
386        self.params = params;
387        self
388    }
389
390    /// Add a single parameter to the group
391    pub fn add_param(mut self, param: Arc<RwLock<Tensor>>) -> Self {
392        self.params.push(param);
393        self
394    }
395
396    /// Set weight decay
397    pub fn weight_decay(mut self, weight_decay: f32) -> Self {
398        self.options
399            .insert("weight_decay".to_string(), weight_decay);
400        self
401    }
402
403    /// Set epsilon
404    pub fn eps(mut self, eps: f32) -> Self {
405        self.options.insert("eps".to_string(), eps);
406        self
407    }
408
409    /// Set a custom option
410    pub fn option(mut self, key: String, value: f32) -> Self {
411        self.options.insert(key, value);
412        self
413    }
414
415    /// Set options from OptimizerOptions
416    pub fn from_options(mut self, options: &OptimizerOptions) -> Self {
417        self.lr = options.lr;
418        self.options = options.to_hashmap();
419        self.options.remove("lr"); // lr is stored separately
420        self
421    }
422
423    /// Build the parameter group
424    pub fn build(self) -> ParamGroup {
425        ParamGroup {
426            params: self.params,
427            lr: self.lr,
428            options: self.options,
429        }
430    }
431}
432
433impl ParamGroup {
434    pub fn new(params: Vec<Arc<RwLock<Tensor>>>, lr: f32) -> Self {
435        Self {
436            params,
437            lr,
438            options: HashMap::new(),
439        }
440    }
441
442    pub fn with_options(mut self, options: HashMap<String, f32>) -> Self {
443        self.options = options;
444        self
445    }
446
447    /// Add a single parameter to the group
448    pub fn add_param(&mut self, param: Arc<RwLock<Tensor>>) {
449        self.params.push(param);
450    }
451
452    /// Get a specific option value, falling back to a default
453    pub fn get_option(&self, key: &str, default: f32) -> f32 {
454        self.options.get(key).copied().unwrap_or(default)
455    }
456
457    /// Set a specific option value
458    pub fn set_option(&mut self, key: String, value: f32) {
459        self.options.insert(key, value);
460    }
461
462    /// Get the number of parameters in this group
463    pub fn param_count(&self) -> usize {
464        self.params.len()
465    }
466
467    /// Check if this group has any parameters
468    pub fn is_empty(&self) -> bool {
469        self.params.is_empty()
470    }
471
472    /// Get all parameters that have gradients
473    pub fn params_with_grads(&self) -> Vec<&Arc<RwLock<Tensor>>> {
474        self.params
475            .iter()
476            .filter(|param| param.read().has_grad())
477            .collect()
478    }
479
480    /// Validate that all parameters in the group are valid
481    pub fn validate(&self) -> bool {
482        !self.params.is_empty() && self.lr.is_finite() && self.lr > 0.0
483    }
484
485    /// Get parameter count for each unique shape in the group
486    pub fn get_shape_counts(&self) -> HashMap<Vec<usize>, usize> {
487        let mut shape_counts = HashMap::new();
488        for param in &self.params {
489            let shape = param.read().shape().dims().to_vec();
490            *shape_counts.entry(shape).or_insert(0) += 1;
491        }
492        shape_counts
493    }
494
495    /// Get total number of parameters (not tensors, but individual parameters)
496    pub fn total_param_count(&self) -> usize {
497        self.params.iter().map(|param| param.read().numel()).sum()
498    }
499
500    /// Clear gradients for all parameters in this group
501    pub fn zero_grad(&self) {
502        for param in &self.params {
503            param.write().zero_grad();
504        }
505    }
506
507    /// Check if any parameter in the group has gradients
508    pub fn has_any_grads(&self) -> bool {
509        self.params.iter().any(|param| param.read().has_grad())
510    }
511
512    /// Get gradient norm for all parameters in the group
513    pub fn grad_norm(&self) -> Result<f32> {
514        let mut total_norm_sq = 0.0f32;
515
516        for param in &self.params {
517            let param_guard = param.read();
518            if let Some(grad) = param_guard.grad() {
519                let grad_norm = grad.norm().map_err(|e| {
520                    TorshError::Other(format!("Failed to compute gradient norm: {}", e))
521                })?;
522                let norm_value = grad_norm.to_vec().map_err(|e| {
523                    TorshError::Other(format!("Failed to extract norm value: {}", e))
524                })?[0];
525                total_norm_sq += norm_value * norm_value;
526            }
527        }
528
529        Ok(total_norm_sq.sqrt())
530    }
531
532    /// Apply gradient clipping to all parameters in the group
533    pub fn clip_grads(&self, max_norm: f32) -> Result<f32> {
534        let total_norm = self.grad_norm()?;
535
536        if total_norm > max_norm {
537            let scale = max_norm / total_norm;
538            for param in &self.params {
539                let mut param_guard = param.write();
540                if let Some(grad) = param_guard.grad() {
541                    let clipped_grad = grad.mul_scalar(scale).map_err(|e| {
542                        TorshError::Other(format!("Failed to clip gradient: {}", e))
543                    })?;
544                    param_guard.set_grad(Some(clipped_grad));
545                }
546            }
547        }
548
549        Ok(total_norm)
550    }
551}
552
553/// Common optimizer options
554#[derive(Debug, Clone)]
555pub struct OptimizerOptions {
556    pub lr: f32,
557    pub weight_decay: f32,
558    pub eps: f32,
559    pub maximize: bool,
560}
561
562impl Default for OptimizerOptions {
563    fn default() -> Self {
564        Self {
565            lr: 1e-3,
566            weight_decay: 0.0,
567            eps: 1e-8,
568            maximize: false,
569        }
570    }
571}
572
573impl OptimizerOptions {
574    /// Create new optimizer options with specified learning rate
575    pub fn new(lr: f32) -> Self {
576        Self {
577            lr,
578            ..Default::default()
579        }
580    }
581
582    /// Set weight decay
583    pub fn with_weight_decay(mut self, weight_decay: f32) -> Self {
584        self.weight_decay = weight_decay;
585        self
586    }
587
588    /// Set epsilon value for numerical stability
589    pub fn with_eps(mut self, eps: f32) -> Self {
590        self.eps = eps;
591        self
592    }
593
594    /// Set maximize flag (for maximization problems)
595    pub fn with_maximize(mut self, maximize: bool) -> Self {
596        self.maximize = maximize;
597        self
598    }
599
600    /// Convert to HashMap for compatibility with parameter groups
601    pub fn to_hashmap(&self) -> HashMap<String, f32> {
602        let mut map = HashMap::new();
603        map.insert("lr".to_string(), self.lr);
604        map.insert("weight_decay".to_string(), self.weight_decay);
605        map.insert("eps".to_string(), self.eps);
606        map.insert(
607            "maximize".to_string(),
608            if self.maximize { 1.0 } else { 0.0 },
609        );
610        map
611    }
612
613    /// Create from HashMap
614    pub fn from_hashmap(map: &HashMap<String, f32>) -> Self {
615        Self {
616            lr: map.get("lr").copied().unwrap_or(1e-3),
617            weight_decay: map.get("weight_decay").copied().unwrap_or(0.0),
618            eps: map.get("eps").copied().unwrap_or(1e-8),
619            maximize: map.get("maximize").copied().unwrap_or(0.0) > 0.0,
620        }
621    }
622
623    /// Validate the options are reasonable
624    pub fn validate(&self) -> Result<()> {
625        if !self.lr.is_finite() || self.lr <= 0.0 {
626            return Err(TorshError::InvalidArgument(
627                "Learning rate must be positive and finite".to_string(),
628            ));
629        }
630        if !self.weight_decay.is_finite() || self.weight_decay < 0.0 {
631            return Err(TorshError::InvalidArgument(
632                "Weight decay must be non-negative and finite".to_string(),
633            ));
634        }
635        if !self.eps.is_finite() || self.eps <= 0.0 {
636            return Err(TorshError::InvalidArgument(
637                "Epsilon must be positive and finite".to_string(),
638            ));
639        }
640        Ok(())
641    }
642
643    /// Create standardized state dict for any optimizer
644    pub fn create_standard_state_dict(
645        optimizer_type: &str,
646        version: Option<&str>,
647        param_groups: &[ParamGroup],
648        state: &HashMap<String, HashMap<String, Tensor>>,
649        global_state: Option<HashMap<String, f32>>,
650    ) -> OptimizerState {
651        let param_group_states = param_groups
652            .iter()
653            .map(|g| ParamGroupState::from_param_group(g))
654            .collect();
655
656        let mut optimizer_state = OptimizerState {
657            optimizer_type: optimizer_type.to_string(),
658            version: version.unwrap_or("1.0").to_string(),
659            param_groups: param_group_states,
660            state: state.clone(),
661            global_state: global_state.unwrap_or_default(),
662        };
663
664        optimizer_state
665    }
666
667    /// Validate state dict compatibility between optimizers
668    pub fn validate_state_compatibility(
669        current_groups: &[ParamGroup],
670        state_groups: &[ParamGroupState],
671    ) -> Result<()> {
672        if current_groups.len() != state_groups.len() {
673            return Err(TorshError::InvalidArgument(format!(
674                "Parameter group count mismatch: expected {}, got {}",
675                current_groups.len(),
676                state_groups.len()
677            )));
678        }
679
680        for (i, (current_group, state_group)) in
681            current_groups.iter().zip(state_groups.iter()).enumerate()
682        {
683            if current_group.params.len() != state_group.param_count {
684                return Err(TorshError::InvalidArgument(format!(
685                    "Parameter count mismatch in group {}: expected {}, got {}",
686                    i,
687                    current_group.params.len(),
688                    state_group.param_count
689                )));
690            }
691        }
692
693        Ok(())
694    }
695}
696
697/// Prelude module for convenient imports
698/// Convergence testing utilities
699#[cfg(test)]
700pub mod convergence_tests {
701    use super::*;
702    use parking_lot::RwLock;
703    use std::ops::Add;
704    use std::sync::Arc;
705    use torsh_tensor::{
706        creation::{randn, zeros},
707        Tensor,
708    };
709
710    /// Test that an optimizer can minimize a simple quadratic function
711    pub fn test_quadratic_convergence<O: Optimizer>(
712        create_optimizer: impl Fn(Vec<Arc<RwLock<Tensor>>>) -> O,
713        tolerance: f32,
714        max_iterations: usize,
715    ) -> Result<()> {
716        // Create a simple quadratic function: f(x) = x^2 + y^2
717        let x = Arc::new(RwLock::new(Tensor::scalar(2.0)?));
718        let y = Arc::new(RwLock::new(Tensor::scalar(2.0)?));
719        let params = vec![x.clone(), y.clone()];
720
721        let mut optimizer = create_optimizer(params);
722
723        for i in 0..max_iterations {
724            // Compute gradients: df/dx = 2x, df/dy = 2y
725            {
726                let x_val = x.read().clone();
727                let y_val = y.read().clone();
728
729                let x_grad = x_val.mul_scalar(2.0)?;
730                let y_grad = y_val.mul_scalar(2.0)?;
731
732                x.write().set_grad(Some(x_grad));
733                y.write().set_grad(Some(y_grad));
734            }
735
736            // Optimizer step
737            optimizer
738                .step()
739                .map_err(|e| TorshError::Other(format!("Optimizer step failed: {}", e)))?;
740
741            // Check convergence
742            let x_val = x.read().to_vec()?[0];
743            let y_val = y.read().to_vec()?[0];
744            let loss = x_val * x_val + y_val * y_val;
745
746            if loss < tolerance {
747                return Ok(());
748            }
749
750            // Clear gradients for next iteration
751            optimizer.zero_grad();
752        }
753
754        Err(TorshError::Other(format!(
755            "Failed to converge within {} iterations",
756            max_iterations
757        )))
758    }
759
760    /// Test that an optimizer can minimize a linear regression problem
761    pub fn test_linear_regression_convergence<O: Optimizer>(
762        create_optimizer: impl Fn(Vec<Arc<RwLock<Tensor>>>) -> O,
763        tolerance: f32,
764        max_iterations: usize,
765    ) -> Result<()> {
766        // Create a simple linear regression problem: y = 2x + 1 + noise
767        let true_weight = 2.0;
768        let true_bias = 1.0;
769
770        // Generate synthetic data
771        let n_samples = 100;
772        let x_data = randn::<f32>(&[n_samples, 1])?;
773        let noise = randn::<f32>(&[n_samples, 1])?.mul_scalar(0.1)?;
774        let y_data = x_data
775            .mul_scalar(true_weight)?
776            .add_scalar(true_bias)?
777            .add(&noise)?;
778
779        // Initialize parameters
780        let weight = Arc::new(RwLock::new(zeros(&[1, 1])?));
781        let bias = Arc::new(RwLock::new(zeros(&[1])?));
782        let params = vec![weight.clone(), bias.clone()];
783
784        let mut optimizer = create_optimizer(params);
785
786        for i in 0..max_iterations {
787            // Forward pass: y_pred = x * weight + bias
788            let w_val = weight.read().clone();
789            let b_val = bias.read().clone();
790
791            let y_pred = x_data.matmul(&w_val)?.add(&b_val)?;
792
793            // Compute loss: MSE = mean((y_pred - y_true)^2)
794            let diff = y_pred.sub(&y_data)?;
795            let loss_tensor = diff.pow(2.0)?.mean(Some(&[0]), false)?;
796            let loss = loss_tensor.to_vec()?[0];
797
798            // Compute gradients
799            let grad_scale = 2.0 / n_samples as f32;
800            let weight_grad = x_data
801                .transpose(0, 1)?
802                .matmul(&diff)?
803                .mul_scalar(grad_scale)?;
804            let bias_grad = diff.sum()?.mul_scalar(grad_scale)?;
805
806            weight.write().set_grad(Some(weight_grad));
807            bias.write().set_grad(Some(bias_grad));
808
809            // Optimizer step
810            optimizer
811                .step()
812                .map_err(|e| TorshError::Other(format!("Optimizer step failed: {}", e)))?;
813
814            // Check convergence
815            if loss < tolerance {
816                // Verify the learned parameters are close to true values
817                let learned_weight = weight.read().to_vec()?[0];
818                let learned_bias = bias.read().to_vec()?[0];
819
820                if (learned_weight - true_weight).abs() < 0.1
821                    && (learned_bias - true_bias).abs() < 0.1
822                {
823                    return Ok(());
824                }
825            }
826
827            // Clear gradients for next iteration
828            optimizer.zero_grad();
829        }
830
831        Err(TorshError::Other(format!(
832            "Failed to converge within {} iterations",
833            max_iterations
834        )))
835    }
836
837    /// Test that an optimizer maintains consistent behavior across multiple runs
838    pub fn test_optimizer_consistency<O: Optimizer>(
839        create_optimizer: impl Fn(Vec<Arc<RwLock<Tensor>>>) -> O,
840        n_runs: usize,
841        tolerance: f32,
842    ) -> Result<()> {
843        let mut final_values = Vec::new();
844
845        for run in 0..n_runs {
846            let param = Arc::new(RwLock::new(Tensor::scalar(1.0)?));
847            let params = vec![param.clone()];
848            let mut optimizer = create_optimizer(params);
849
850            // Run for a fixed number of steps
851            for _ in 0..10 {
852                {
853                    let param_val = param.read().clone();
854                    let grad = param_val.mul_scalar(2.0)?; // Simple gradient
855                    param.write().set_grad(Some(grad));
856                }
857
858                optimizer
859                    .step()
860                    .map_err(|e| TorshError::Other(format!("Optimizer step failed: {}", e)))?;
861                optimizer.zero_grad();
862            }
863
864            final_values.push(param.read().to_vec()?[0]);
865        }
866
867        // Check that all runs produce similar results
868        let mean_value = final_values.iter().sum::<f32>() / final_values.len() as f32;
869        for &value in &final_values {
870            if (value - mean_value).abs() > tolerance {
871                return Err(TorshError::Other(format!(
872                    "Inconsistent optimizer behavior: values vary by more than {}",
873                    tolerance
874                )));
875            }
876        }
877
878        Ok(())
879    }
880}
881
882pub mod prelude {
883    pub use crate::adabelief::AdaBelief;
884    pub use crate::adabound::AdaBound;
885    pub use crate::adadelta::AdaDelta;
886    pub use crate::adagrad::AdaGrad;
887    pub use crate::adahessian::{AdaHessian, AdaHessianBuilder};
888    pub use crate::adam::{Adam, AdamW};
889    pub use crate::adamax::AdaMax;
890    pub use crate::asgd::ASGD;
891    pub use crate::checkpointing::{
892        Checkpoint, CheckpointConfig, CheckpointManager, CheckpointMetadata, CheckpointStatistics,
893        CheckpointSupport, CheckpointingOptimizer,
894    };
895    pub use crate::composition::{
896        CombinationMethod, ComposedOptimizer, CompositionBuilder, CompositionStrategy,
897        OptimizerMetrics, SwitchCriterion, VotingMethod,
898    };
899    pub use crate::debugging::{
900        AnalysisReport, AnalyzerConfig, ConvergenceTracker, GradientFlowPoint, GradientStatistics,
901        HyperparameterSensitivity, OptimizationRecommendation, OptimizationStep, OptimizerAnalyzer,
902        ParameterStatistics, RecommendationCategory, SensitivityReport, SensitivityResult,
903        Severity,
904    };
905    pub use crate::distributed::{
906        utils as distributed_utils, AsyncConfig, AsyncSGD, CommunicationStats, DistributedBackend,
907        DistributedConfig, DistributedOptimizer, ElasticAveragingSGD, SyncStrategy,
908    };
909    pub use crate::ftrl::{FTRLBuilder, FTRL};
910    pub use crate::fused_kernels::{
911        fused_adadelta_step, fused_adagrad_step, fused_adam_step, fused_rmsprop_step,
912        fused_sgd_step, FusedKernelSupport, FusedStats,
913    };
914    pub use crate::grad_accumulation::{
915        with_gradient_accumulation, AccumulatingOptimizer, GradientAccumulationSupport,
916        GradientAccumulator,
917    };
918    pub use crate::kfac::{KFACBuilder, KFAC};
919    pub use crate::lamb::LAMB;
920    pub use crate::lazy_updates::{
921        LazyUpdateConfig, LazyUpdateDecision, LazyUpdateManager, LazyUpdateOptimizer,
922        LazyUpdateStatistics, LazyUpdateSupport, ParameterImportance, PendingUpdate,
923        UpdatePriority,
924    };
925    pub use crate::lbfgs::LBFGS;
926    pub use crate::lion::{Lion, LionBuilder, LionConfig};
927    pub use crate::lookahead::{lookahead_adam, lookahead_radam, lookahead_sgd, Lookahead};
928    pub use crate::low_precision::{
929        LowPrecisionConvertible, LowPrecisionOptimizer, LowPrecisionState, PrecisionType,
930        StateStatistics,
931    };
932    pub use crate::lr_scheduler::{
933        CosineAnnealingLR, ExponentialLR, LRScheduler, OneCycleLR, ReduceLROnPlateau, StepLR,
934    };
935    pub use crate::lr_scheduler_additional::{
936        ConstantLR, CosineAnnealingWarmRestarts, CyclicLR, LinearLR, MultiStepLR, PolynomialLR,
937    };
938    pub use crate::lr_scheduler_enhanced::{
939        utils as lr_enhanced_utils, AdaptiveLRScheduler, AdaptiveSchedulerStats, AdaptiveStrategy,
940        CosineAnnealingWarmRestartsWithWarmup, PolynomialDecayWithWarmup, WarmupStrategy,
941    };
942    pub use crate::memory_efficient::{
943        CircularBuffer, MemoryConfig, MemoryEfficientAdam, MemoryEfficientLBFGS,
944        MemoryEfficientOptimizerBuilder, MemoryPool,
945    };
946    pub use crate::memory_mapped::{
947        MemoryMappedConfig, MemoryMappedFile, MemoryMappedOptimizer, MemoryMappedStateStorage,
948        MemoryMappedSupport, StorageStatistics,
949    };
950    pub use crate::mixed_precision::{
951        with_mixed_precision, MixedPrecisionConfig, MixedPrecisionOptimizer,
952    };
953    pub use crate::nadam::NAdam;
954    pub use crate::natural_gradient::{NaturalGradient, NaturalGradientBuilder};
955    pub use crate::newton_cg::{NewtonCG, NewtonCGBuilder, NewtonCGConfig};
956    pub use crate::online_learning::{
957        OnlineGradientDescent, ProximalGradient, ProximalOperator, SAGA, SVRG,
958    };
959    pub use crate::prodigy::{Prodigy, ProdigyBuilder, ProdigyConfig};
960    pub use crate::radam::RAdam;
961    pub use crate::ranger::{Ranger, RangerBuilder};
962    pub use crate::rmsprop::RMSprop;
963    pub use crate::rprop::Rprop;
964    pub use crate::schedule_free::{ScheduleFreeAdamW, ScheduleFreeAdamWBuilder};
965    pub use crate::sgd::SGD;
966    pub use crate::shampoo::{Shampoo, ShampooBuilder};
967    pub use crate::sophia::{Sophia, SophiaBuilder, SophiaConfig};
968    pub use crate::sparse_adam::SparseAdam;
969    pub use crate::state_dict_ops::{
970        CompressionMethod, CompressionStats, MemoryEstimate, SerializationFormat, StateDictConfig,
971        StateDictManager,
972    };
973    pub use crate::trust_region::{
974        SubproblemSolver, TrustRegionBuilder, TrustRegionConfig, TrustRegionMethod,
975        TrustRegionStrategy,
976    };
977    pub use crate::yellowfin::{YellowFin, YellowFinBuilder, YellowFinConfig};
978    pub use crate::{Optimizer, OptimizerOptions, OptimizerState, ParamGroup, ParamGroupBuilder};
979    pub use crate::{OptimizerError, OptimizerResult};
980}
981
982// Re-export commonly used types
983pub use adam::{Adam, AdamW};
984pub use distributed::{DistributedBackend, DistributedConfig, DistributedOptimizer, SyncStrategy};
985pub use rmsprop::RMSprop;
986pub use sgd::SGD;
987
988#[cfg(test)]
989mod tests {
990    use super::*;
991
992    #[test]
993    fn test_param_group() {
994        let params = vec![];
995        let group = ParamGroup::new(params, 0.01);
996        assert_eq!(group.lr, 0.01);
997    }
998}
torsh_optim/lib.rs

torsh_optim/
lib.rs
