optirs_core/
unified_api.rs

1// Unified API consistent with popular deep learning frameworks
2//
3// This module provides a unified interface that closely follows the design patterns
4// of popular deep learning frameworks like PyTorch, TensorFlow, and JAX/Optax.
5//
6// # Design Principles
7//
8// - **Parameter Groups**: Support for different optimization parameters for different layers
9// - **State Management**: Automatic handling of optimizer state
10// - **Framework Consistency**: APIs that feel familiar to PyTorch/TensorFlow users
11// - **Flexible Configuration**: Easy-to-use builder patterns
12// - **Scheduler Integration**: Seamless integration with learning rate schedulers
13
14use crate::error::{OptimError, Result};
15use crate::schedulers::LearningRateScheduler;
16use scirs2_core::ndarray::{Array, Array1, Dimension, ScalarOperand};
17use scirs2_core::numeric::Float;
18use serde::{Deserialize, Serialize};
19use std::collections::HashMap;
20use std::fmt::Debug;
21
22/// Unified optimizer configuration
23#[derive(Debug, Clone, Serialize, Deserialize)]
24pub struct OptimizerConfig<A: Float> {
25    /// Learning rate
26    pub lr: A,
27    /// Weight decay (L2 regularization)
28    pub weight_decay: A,
29    /// Gradient clipping value (optional)
30    pub grad_clip: Option<A>,
31    /// Additional optimizer-specific parameters
32    pub params: HashMap<String, A>,
33}
34
35impl<A: Float + Send + Sync> Default for OptimizerConfig<A> {
36    fn default() -> Self {
37        Self {
38            lr: A::from(0.001)
39                .expect("OptimizerConfig: default learning rate (0.001) must fit in A"),
40            weight_decay: A::zero(),
41            grad_clip: None,
42            params: HashMap::new(),
43        }
44    }
45}
46
47impl<A: Float + Send + Sync> OptimizerConfig<A> {
48    /// Create a new optimizer configuration with the given learning rate
49    pub fn new(lr: A) -> Self {
50        Self {
51            lr,
52            ..Default::default()
53        }
54    }
55
56    /// Set weight decay
57    pub fn weight_decay(mut self, weightdecay: A) -> Self {
58        self.weight_decay = weightdecay;
59        self
60    }
61
62    /// Set gradient clipping
63    pub fn grad_clip(mut self, gradclip: A) -> Self {
64        self.grad_clip = Some(gradclip);
65        self
66    }
67
68    /// Add a custom parameter
69    pub fn param<S: Into<String>>(mut self, key: S, value: A) -> Self {
70        self.params.insert(key.into(), value);
71        self
72    }
73
74    /// Set multiple parameters at once
75    pub fn params(mut self, params: HashMap<String, A>) -> Self {
76        self.params.extend(params);
77        self
78    }
79}
80
81/// Parameter tensor wrapper for unified API
82#[derive(Debug, Clone)]
83pub struct Parameter<A: Float, D: Dimension> {
84    /// Parameter data
85    pub data: Array<A, D>,
86    /// Gradient data (optional)
87    pub grad: Option<Array<A, D>>,
88    /// Whether this parameter requires gradients
89    pub requires_grad: bool,
90    /// Parameter name/identifier
91    pub name: String,
92}
93
94impl<A: Float + ScalarOperand, D: Dimension + Send + Sync> Parameter<A, D> {
95    /// Create a new parameter
96    pub fn new<S: Into<String>>(data: Array<A, D>, name: S) -> Self {
97        Self {
98            data,
99            grad: None,
100            requires_grad: true,
101            name: name.into(),
102        }
103    }
104
105    /// Create a parameter that doesn't require gradients
106    pub fn no_grad<S: Into<String>>(data: Array<A, D>, name: S) -> Self {
107        Self {
108            data,
109            grad: None,
110            requires_grad: false,
111            name: name.into(),
112        }
113    }
114
115    /// Set gradient for this parameter
116    pub fn set_grad(&mut self, grad: Array<A, D>) {
117        if self.requires_grad {
118            self.grad = Some(grad);
119        }
120    }
121
122    /// Clear gradients
123    pub fn zero_grad(&mut self) {
124        self.grad = None;
125    }
126
127    /// Get gradient reference
128    pub fn grad(&self) -> Option<&Array<A, D>> {
129        self.grad.as_ref()
130    }
131
132    /// Apply gradient clipping if specified
133    pub fn clip_grad(&mut self, maxnorm: A) -> Result<()> {
134        if let Some(ref mut grad) = self.grad {
135            let _norm = grad
136                .iter()
137                .map(|x| (*x) * (*x))
138                .fold(A::zero(), |acc, x| acc + x)
139                .sqrt();
140            if _norm > maxnorm {
141                let scale = maxnorm / _norm;
142                grad.mapv_inplace(|x| x * scale);
143            }
144        }
145        Ok(())
146    }
147}
148
149/// Unified optimizer interface
150pub trait UnifiedOptimizer<A: Float> {
151    /// Get optimizer configuration
152    fn config(&self) -> &OptimizerConfig<A>;
153
154    /// Update a single parameter
155    fn step_param<D: Dimension>(&mut self, param: &mut Parameter<A, D>) -> Result<()>
156    where
157        A: ScalarOperand + Debug;
158
159    /// Update multiple parameters
160    fn step_params<D: Dimension>(&mut self, params: &mut [Parameter<A, D>]) -> Result<()>
161    where
162        A: ScalarOperand + Debug,
163    {
164        for param in params.iter_mut() {
165            self.step_param(param)?;
166        }
167        Ok(())
168    }
169
170    /// Zero gradients for all parameters
171    fn zero_grad<D: Dimension>(&self, params: &mut [Parameter<A, D>]) {
172        for param in params.iter_mut() {
173            param.grad = None;
174        }
175    }
176
177    /// Update learning rate
178    fn set_lr(&mut self, lr: A);
179
180    /// Get current learning rate
181    fn get_lr(&self) -> A;
182
183    /// Serialize the full optimizer state (configuration + per-parameter buffers)
184    ///
185    /// Values are stored as little-endian `f64` blobs, so the format is dependency
186    /// free and stable across `f32` / `f64` optimizers. Keys are namespaced:
187    ///
188    /// | key | contents |
189    /// |-----|----------|
190    /// | `format.version` | one `f64`, currently `1` |
191    /// | `config.lr` / `config.weight_decay` | one `f64` each |
192    /// | `config.grad_clip` | one `f64`, absent when clipping is disabled |
193    /// | `config.param.<name>` | one `f64` per optimizer-specific hyperparameter |
194    /// | `<buffer>.<parameter name>` | the buffer contents, one `f64` per element |
195    fn state_dict(&self) -> Result<HashMap<String, Vec<u8>>>;
196
197    /// Restore state previously produced by [`UnifiedOptimizer::state_dict`]
198    ///
199    /// # Errors
200    ///
201    /// Returns an error when the payload is truncated (not a whole number of `f64`
202    /// values), when the format version is unknown, or when a restored buffer does
203    /// not match the shape of the buffer it replaces.
204    fn load_state_dict(&mut self, statedict: HashMap<String, Vec<u8>>) -> Result<()>;
205}
206
207/// Serialization format version written into every state dictionary
208const STATE_DICT_VERSION: f64 = 1.0;
209
210/// Key holding the state-dictionary format version
211const KEY_FORMAT_VERSION: &str = "format.version";
212
213/// Encodes `f64` values as a little-endian byte blob
214fn encode_f64_slice(values: &[f64]) -> Vec<u8> {
215    let mut bytes = Vec::with_capacity(values.len() * 8);
216    for value in values {
217        bytes.extend_from_slice(&value.to_le_bytes());
218    }
219    bytes
220}
221
222/// Decodes a little-endian byte blob back into `f64` values
223fn decode_f64_slice(key: &str, bytes: &[u8]) -> Result<Vec<f64>> {
224    if !bytes.len().is_multiple_of(8) {
225        return Err(OptimError::InvalidConfig(format!(
226            "state dict entry '{}' is truncated: {} bytes is not a multiple of 8",
227            key,
228            bytes.len()
229        )));
230    }
231
232    let mut values = Vec::with_capacity(bytes.len() / 8);
233    for chunk in bytes.chunks_exact(8) {
234        let mut buf = [0u8; 8];
235        buf.copy_from_slice(chunk);
236        values.push(f64::from_le_bytes(buf));
237    }
238    Ok(values)
239}
240
241/// Converts a floating-point value into the state-dict representation
242fn to_state_value<A: Float>(key: &str, value: A) -> Result<f64> {
243    value.to_f64().ok_or_else(|| {
244        OptimError::InvalidConfig(format!("state dict entry '{}' is not representable", key))
245    })
246}
247
248/// Converts a state-dict value back into the optimizer's floating-point type
249fn from_state_value<A: Float>(key: &str, value: f64) -> Result<A> {
250    A::from(value).ok_or_else(|| {
251        OptimError::InvalidConfig(format!(
252            "state dict entry '{}' holds a value that is not representable in the target type",
253            key
254        ))
255    })
256}
257
258/// Reads exactly one scalar out of a state-dict entry
259fn decode_scalar(key: &str, bytes: &[u8]) -> Result<f64> {
260    let values = decode_f64_slice(key, bytes)?;
261    match values.as_slice() {
262        [single] => Ok(*single),
263        other => Err(OptimError::InvalidConfig(format!(
264            "state dict entry '{}' must hold exactly one value, found {}",
265            key,
266            other.len()
267        ))),
268    }
269}
270
271/// Serializes the shared [`OptimizerConfig`] portion of a state dictionary
272fn encode_config<A: Float>(
273    config: &OptimizerConfig<A>,
274    target: &mut HashMap<String, Vec<u8>>,
275) -> Result<()> {
276    target.insert(
277        KEY_FORMAT_VERSION.to_string(),
278        encode_f64_slice(&[STATE_DICT_VERSION]),
279    );
280    target.insert(
281        "config.lr".to_string(),
282        encode_f64_slice(&[to_state_value("config.lr", config.lr)?]),
283    );
284    target.insert(
285        "config.weight_decay".to_string(),
286        encode_f64_slice(&[to_state_value("config.weight_decay", config.weight_decay)?]),
287    );
288    if let Some(clip) = config.grad_clip {
289        target.insert(
290            "config.grad_clip".to_string(),
291            encode_f64_slice(&[to_state_value("config.grad_clip", clip)?]),
292        );
293    }
294    for (name, value) in config.params.iter() {
295        let key = format!("config.param.{}", name);
296        let encoded = encode_f64_slice(&[to_state_value(&key, *value)?]);
297        target.insert(key, encoded);
298    }
299    Ok(())
300}
301
302/// Restores the shared [`OptimizerConfig`] portion of a state dictionary
303fn decode_config<A: Float + Send + Sync>(
304    state: &HashMap<String, Vec<u8>>,
305    config: &mut OptimizerConfig<A>,
306) -> Result<()> {
307    let version_bytes = state.get(KEY_FORMAT_VERSION).ok_or_else(|| {
308        OptimError::InvalidConfig(format!("state dict is missing '{}'", KEY_FORMAT_VERSION))
309    })?;
310    let version = decode_scalar(KEY_FORMAT_VERSION, version_bytes)?;
311    if version != STATE_DICT_VERSION {
312        return Err(OptimError::InvalidConfig(format!(
313            "unsupported state dict version {} (expected {})",
314            version, STATE_DICT_VERSION
315        )));
316    }
317
318    if let Some(bytes) = state.get("config.lr") {
319        config.lr = from_state_value("config.lr", decode_scalar("config.lr", bytes)?)?;
320    }
321    if let Some(bytes) = state.get("config.weight_decay") {
322        config.weight_decay = from_state_value(
323            "config.weight_decay",
324            decode_scalar("config.weight_decay", bytes)?,
325        )?;
326    }
327    config.grad_clip = match state.get("config.grad_clip") {
328        Some(bytes) => Some(from_state_value(
329            "config.grad_clip",
330            decode_scalar("config.grad_clip", bytes)?,
331        )?),
332        None => None,
333    };
334
335    for (key, bytes) in state.iter() {
336        if let Some(name) = key.strip_prefix("config.param.") {
337            let value = from_state_value(key, decode_scalar(key, bytes)?)?;
338            config.params.insert(name.to_string(), value);
339        }
340    }
341    Ok(())
342}
343
344/// Restores a named collection of `Array1` buffers, validating their shapes
345fn decode_buffers<A: Float>(
346    state: &HashMap<String, Vec<u8>>,
347    prefix: &str,
348    existing: &HashMap<String, Array1<A>>,
349) -> Result<HashMap<String, Array1<A>>> {
350    let mut restored = HashMap::new();
351    for (key, bytes) in state.iter() {
352        let name = match key.strip_prefix(prefix) {
353            Some(name) => name,
354            None => continue,
355        };
356        let values = decode_f64_slice(key, bytes)?;
357
358        if let Some(current) = existing.get(name) {
359            if current.len() != values.len() {
360                return Err(OptimError::DimensionMismatch(format!(
361                    "state dict buffer '{}' has {} elements but the optimizer holds {}",
362                    key,
363                    values.len(),
364                    current.len()
365                )));
366            }
367        }
368
369        let mut buffer = Array1::zeros(values.len());
370        for (slot, value) in buffer.iter_mut().zip(values.iter()) {
371            *slot = from_state_value(key, *value)?;
372        }
373        restored.insert(name.to_string(), buffer);
374    }
375    Ok(restored)
376}
377
378/// Serializes a named collection of `Array1` buffers
379fn encode_buffers<A: Float>(
380    buffers: &HashMap<String, Array1<A>>,
381    prefix: &str,
382    target: &mut HashMap<String, Vec<u8>>,
383) -> Result<()> {
384    for (name, buffer) in buffers.iter() {
385        let key = format!("{}{}", prefix, name);
386        let mut values = Vec::with_capacity(buffer.len());
387        for value in buffer.iter() {
388            values.push(to_state_value(&key, *value)?);
389        }
390        target.insert(key, encode_f64_slice(&values));
391    }
392    Ok(())
393}
394
395/// SGD optimizer with unified API
396#[derive(Debug)]
397pub struct UnifiedSGD<A: Float> {
398    config: OptimizerConfig<A>,
399    momentum_buffers: HashMap<String, Array1<A>>,
400}
401
402impl<A: Float + ScalarOperand + Debug + Send + Sync> UnifiedSGD<A> {
403    /// Create a new SGD optimizer
404    pub fn new(config: OptimizerConfig<A>) -> Self {
405        Self {
406            config,
407            momentum_buffers: HashMap::new(),
408        }
409    }
410
411    /// Create SGD with momentum
412    pub fn with_momentum(mut config: OptimizerConfig<A>, momentum: A) -> Self {
413        config.params.insert("momentum".to_string(), momentum);
414        Self::new(config)
415    }
416}
417
418impl<A: Float + ScalarOperand + Debug + Send + Sync> UnifiedOptimizer<A> for UnifiedSGD<A> {
419    fn config(&self) -> &OptimizerConfig<A> {
420        &self.config
421    }
422
423    fn step_param<D: Dimension>(&mut self, param: &mut Parameter<A, D>) -> Result<()> {
424        if !param.requires_grad {
425            return Ok(());
426        }
427
428        // Check gradient exists first
429        if param.grad.is_none() {
430            return Err(OptimError::InvalidConfig(
431                "Parameter has no gradient".to_string(),
432            ));
433        }
434
435        // Apply gradient clipping if configured
436        if let Some(max_norm) = self.config.grad_clip {
437            param.clip_grad(max_norm)?;
438        }
439
440        // Apply weight decay
441        if self.config.weight_decay > A::zero() {
442            param
443                .data
444                .mapv_inplace(|x| x * (A::one() - self.config.weight_decay * self.config.lr));
445        }
446
447        // Get gradient safely (guaranteed `Some` by the `is_none()` guard above)
448        let grad = param
449            .grad
450            .as_ref()
451            .ok_or_else(|| OptimError::InvalidConfig("Parameter has no gradient".to_string()))?;
452
453        // Get momentum factor
454        let momentum = self
455            .config
456            .params
457            .get("momentum")
458            .copied()
459            .unwrap_or(A::zero());
460
461        if momentum > A::zero() {
462            // SGD with momentum
463            if let Some(momentum_buffer) = self.momentum_buffers.get_mut(&param.name) {
464                // Update momentum buffer
465                for (m, g) in momentum_buffer.iter_mut().zip(grad.iter()) {
466                    *m = momentum * (*m) + *g;
467                }
468                // Update parameters
469                for (p, m) in param.data.iter_mut().zip(momentum_buffer.iter()) {
470                    *p = *p - self.config.lr * (*m);
471                }
472            } else {
473                // Initialize momentum buffer
474                let mut momentum_buffer = Array1::zeros(grad.len());
475                for (m, g) in momentum_buffer.iter_mut().zip(grad.iter()) {
476                    *m = *g;
477                }
478                // Update parameters
479                for (p, m) in param.data.iter_mut().zip(momentum_buffer.iter()) {
480                    *p = *p - self.config.lr * (*m);
481                }
482                self.momentum_buffers
483                    .insert(param.name.clone(), momentum_buffer);
484            }
485        } else {
486            // Standard SGD
487            for (p, g) in param.data.iter_mut().zip(grad.iter()) {
488                *p = *p - self.config.lr * (*g);
489            }
490        }
491
492        Ok(())
493    }
494
495    fn set_lr(&mut self, lr: A) {
496        self.config.lr = lr;
497    }
498
499    fn get_lr(&self) -> A {
500        self.config.lr
501    }
502
503    fn state_dict(&self) -> Result<HashMap<String, Vec<u8>>> {
504        let mut state = HashMap::new();
505        encode_config(&self.config, &mut state)?;
506        encode_buffers(&self.momentum_buffers, "sgd.momentum_buffer.", &mut state)?;
507        Ok(state)
508    }
509
510    fn load_state_dict(&mut self, statedict: HashMap<String, Vec<u8>>) -> Result<()> {
511        decode_config(&statedict, &mut self.config)?;
512        self.momentum_buffers =
513            decode_buffers(&statedict, "sgd.momentum_buffer.", &self.momentum_buffers)?;
514        Ok(())
515    }
516}
517
518/// Adam optimizer with unified API
519#[derive(Debug)]
520pub struct UnifiedAdam<A: Float> {
521    config: OptimizerConfig<A>,
522    /// Per-parameter update counters driving bias correction
523    ///
524    /// Adam's bias correction depends on how many updates *that particular tensor*
525    /// has received. A single shared counter would advance once per parameter in the
526    /// model, so a 100-tensor model would reach t = 100 after a single optimizer
527    /// step and its bias correction would be wrong for every tensor.
528    step_counts: HashMap<String, usize>,
529    exp_avg: HashMap<String, Array1<A>>,
530    exp_avg_sq: HashMap<String, Array1<A>>,
531}
532
533impl<A: Float + ScalarOperand + Debug + Send + Sync> UnifiedAdam<A> {
534    /// Create a new Adam optimizer
535    pub fn new(config: OptimizerConfig<A>) -> Self {
536        let mut params = config.params.clone();
537        params.entry("beta1".to_string()).or_insert_with(|| {
538            A::from(0.9).expect("UnifiedAdam: default beta1 (0.9) must fit in A")
539        });
540        params.entry("beta2".to_string()).or_insert_with(|| {
541            A::from(0.999).expect("UnifiedAdam: default beta2 (0.999) must fit in A")
542        });
543        params.entry("eps".to_string()).or_insert_with(|| {
544            A::from(1e-8).expect("UnifiedAdam: default eps (1e-8) must fit in A")
545        });
546
547        Self {
548            config: OptimizerConfig { params, ..config },
549            step_counts: HashMap::new(),
550            exp_avg: HashMap::new(),
551            exp_avg_sq: HashMap::new(),
552        }
553    }
554
555    /// Number of updates applied to the parameter called `name`
556    pub fn step_count(&self, name: &str) -> usize {
557        self.step_counts.get(name).copied().unwrap_or(0)
558    }
559
560    /// Create Adam with custom betas
561    pub fn with_betas(mut config: OptimizerConfig<A>, beta1: A, beta2: A) -> Self {
562        config.params.insert("beta1".to_string(), beta1);
563        config.params.insert("beta2".to_string(), beta2);
564        Self::new(config)
565    }
566}
567
568impl<A: Float + ScalarOperand + Debug + Send + Sync> UnifiedOptimizer<A> for UnifiedAdam<A> {
569    fn config(&self) -> &OptimizerConfig<A> {
570        &self.config
571    }
572
573    fn step_param<D: Dimension>(&mut self, param: &mut Parameter<A, D>) -> Result<()> {
574        if !param.requires_grad {
575            return Ok(());
576        }
577
578        // Check gradient exists first
579        if param.grad.is_none() {
580            return Err(OptimError::InvalidConfig(
581                "Parameter has no gradient".to_string(),
582            ));
583        }
584
585        // Apply gradient clipping if configured
586        if let Some(max_norm) = self.config.grad_clip {
587            param.clip_grad(max_norm)?;
588        }
589
590        // Advance this parameter's own clock, not a counter shared by every tensor.
591        let step_count = {
592            let counter = self.step_counts.entry(param.name.clone()).or_insert(0);
593            *counter = counter.saturating_add(1);
594            *counter
595        };
596
597        let beta1 = *self.config.params.get("beta1").ok_or_else(|| {
598            OptimError::InvalidConfig("Adam configuration is missing 'beta1'".to_string())
599        })?;
600        let beta2 = *self.config.params.get("beta2").ok_or_else(|| {
601            OptimError::InvalidConfig("Adam configuration is missing 'beta2'".to_string())
602        })?;
603        let eps = *self.config.params.get("eps").ok_or_else(|| {
604            OptimError::InvalidConfig("Adam configuration is missing 'eps'".to_string())
605        })?;
606
607        // Get gradient safely
608        let grad = param
609            .grad
610            .as_ref()
611            .ok_or_else(|| OptimError::InvalidConfig("Parameter has no gradient".to_string()))?;
612
613        // Initialize or get existing moment estimates
614        let exp_avg = self
615            .exp_avg
616            .entry(param.name.clone())
617            .or_insert_with(|| Array1::zeros(grad.len()));
618        let exp_avg_sq = self
619            .exp_avg_sq
620            .entry(param.name.clone())
621            .or_insert_with(|| Array1::zeros(grad.len()));
622
623        // Update biased first and second moment estimates
624        for ((exp_avg_val, exp_avg_sq_val), grad_val) in exp_avg
625            .iter_mut()
626            .zip(exp_avg_sq.iter_mut())
627            .zip(grad.iter())
628        {
629            *exp_avg_val = beta1 * (*exp_avg_val) + (A::one() - beta1) * (*grad_val);
630            *exp_avg_sq_val =
631                beta2 * (*exp_avg_sq_val) + (A::one() - beta2) * (*grad_val) * (*grad_val);
632        }
633
634        // Bias correction driven by this parameter's own step count
635        let exponent = i32::try_from(step_count).map_err(|_| {
636            OptimError::InvalidConfig(
637                "Timestep too large for bias correction calculation".to_string(),
638            )
639        })?;
640        let bias_correction1 = A::one() - beta1.powi(exponent);
641        let bias_correction2 = A::one() - beta2.powi(exponent);
642
643        let step_size = self.config.lr * (bias_correction2.sqrt() / bias_correction1);
644
645        // Update parameters
646        for ((p, exp_avg_val), exp_avg_sq_val) in param
647            .data
648            .iter_mut()
649            .zip(exp_avg.iter())
650            .zip(exp_avg_sq.iter())
651        {
652            let denom = exp_avg_sq_val.sqrt() + eps;
653            *p = *p - step_size * (*exp_avg_val) / denom;
654        }
655
656        // Apply weight decay after the main update
657        if self.config.weight_decay > A::zero() {
658            param
659                .data
660                .mapv_inplace(|x| x * (A::one() - self.config.weight_decay * self.config.lr));
661        }
662
663        Ok(())
664    }
665
666    fn set_lr(&mut self, lr: A) {
667        self.config.lr = lr;
668    }
669
670    fn get_lr(&self) -> A {
671        self.config.lr
672    }
673
674    fn state_dict(&self) -> Result<HashMap<String, Vec<u8>>> {
675        let mut state = HashMap::new();
676        encode_config(&self.config, &mut state)?;
677        encode_buffers(&self.exp_avg, "adam.exp_avg.", &mut state)?;
678        encode_buffers(&self.exp_avg_sq, "adam.exp_avg_sq.", &mut state)?;
679        for (name, count) in self.step_counts.iter() {
680            state.insert(
681                format!("adam.step_count.{}", name),
682                encode_f64_slice(&[*count as f64]),
683            );
684        }
685        Ok(state)
686    }
687
688    fn load_state_dict(&mut self, statedict: HashMap<String, Vec<u8>>) -> Result<()> {
689        decode_config(&statedict, &mut self.config)?;
690
691        let exp_avg = decode_buffers(&statedict, "adam.exp_avg.", &self.exp_avg)?;
692        let exp_avg_sq = decode_buffers(&statedict, "adam.exp_avg_sq.", &self.exp_avg_sq)?;
693
694        // The two moment buffers describe the same tensors and must agree.
695        for (name, buffer) in exp_avg.iter() {
696            match exp_avg_sq.get(name) {
697                Some(other) if other.len() == buffer.len() => {}
698                Some(other) => {
699                    return Err(OptimError::DimensionMismatch(format!(
700                        "state dict moments for '{}' disagree: {} vs {} elements",
701                        name,
702                        buffer.len(),
703                        other.len()
704                    )))
705                }
706                None => {
707                    return Err(OptimError::InvalidConfig(format!(
708                        "state dict has 'adam.exp_avg.{}' but no matching 'adam.exp_avg_sq' entry",
709                        name
710                    )))
711                }
712            }
713        }
714
715        let mut step_counts = HashMap::new();
716        for (key, bytes) in statedict.iter() {
717            if let Some(name) = key.strip_prefix("adam.step_count.") {
718                let value = decode_scalar(key, bytes)?;
719                if !value.is_finite() || value < 0.0 {
720                    return Err(OptimError::InvalidConfig(format!(
721                        "state dict entry '{}' holds an invalid step count {}",
722                        key, value
723                    )));
724                }
725                step_counts.insert(name.to_string(), value as usize);
726            }
727        }
728
729        self.exp_avg = exp_avg;
730        self.exp_avg_sq = exp_avg_sq;
731        self.step_counts = step_counts;
732        Ok(())
733    }
734}
735
736/// Optimizer factory for creating optimizers with unified API
737pub struct OptimizerFactory;
738
739impl OptimizerFactory {
740    /// Create SGD optimizer
741    pub fn sgd<A: Float + ScalarOperand + Debug + Send + Sync>(
742        config: OptimizerConfig<A>,
743    ) -> UnifiedSGD<A> {
744        UnifiedSGD::new(config)
745    }
746
747    /// Create Adam optimizer
748    pub fn adam<A: Float + ScalarOperand + Debug + Send + Sync>(
749        config: OptimizerConfig<A>,
750    ) -> UnifiedAdam<A> {
751        UnifiedAdam::new(config)
752    }
753
754    /// Create SGD with momentum
755    pub fn sgd_momentum<A: Float + ScalarOperand + Debug + Send + Sync>(
756        config: OptimizerConfig<A>,
757        momentum: A,
758    ) -> UnifiedSGD<A> {
759        UnifiedSGD::with_momentum(config, momentum)
760    }
761
762    /// Create Adam with custom parameters
763    pub fn adam_custom<A: Float + ScalarOperand + Debug + Send + Sync>(
764        config: OptimizerConfig<A>,
765        beta1: A,
766        beta2: A,
767    ) -> UnifiedAdam<A> {
768        UnifiedAdam::with_betas(config, beta1, beta2)
769    }
770}
771
772/// Training loop helper with unified API
773pub struct TrainingLoop<A: Float, O: UnifiedOptimizer<A>> {
774    optimizer: O,
775    scheduler: Option<Box<dyn LearningRateScheduler<A>>>,
776    _phantom: std::marker::PhantomData<A>,
777}
778
779impl<A: Float + ScalarOperand + Debug, O: UnifiedOptimizer<A> + Send + Sync> TrainingLoop<A, O> {
780    /// Create a new training loop
781    pub fn new(optimizer: O) -> Self {
782        Self {
783            optimizer,
784            scheduler: None,
785            _phantom: std::marker::PhantomData,
786        }
787    }
788
789    /// Add a learning rate scheduler
790    pub fn with_scheduler(mut self, scheduler: Box<dyn LearningRateScheduler<A>>) -> Self {
791        self.scheduler = Some(scheduler);
792        self
793    }
794
795    /// Perform one training step
796    pub fn step<D: Dimension>(&mut self, params: &mut [Parameter<A, D>]) -> Result<()> {
797        // Update parameters
798        self.optimizer.step_params(params)?;
799
800        // Update learning rate if scheduler is present
801        if let Some(ref mut scheduler) = self.scheduler {
802            let new_lr = scheduler.step();
803            self.optimizer.set_lr(new_lr);
804        }
805
806        Ok(())
807    }
808
809    /// Zero gradients
810    pub fn zero_grad<D: Dimension>(&self, params: &mut [Parameter<A, D>]) {
811        for param in params.iter_mut() {
812            param.grad = None;
813        }
814    }
815
816    /// Get current learning rate
817    pub fn get_lr(&self) -> A {
818        self.optimizer.get_lr()
819    }
820}
821
822#[cfg(test)]
823mod tests {
824    use super::*;
825    use scirs2_core::ndarray::Array1;
826
827    #[test]
828    fn test_unified_sgd() {
829        let config = OptimizerConfig::new(0.1f64);
830        let mut optimizer = UnifiedSGD::new(config);
831
832        let mut param = Parameter::new(Array1::from_vec(vec![1.0, 2.0, 3.0]), "test_param");
833        param.set_grad(Array1::from_vec(vec![0.1, 0.2, 0.3]));
834
835        optimizer
836            .step_param(&mut param)
837            .expect("optimizer.step_param succeeds in test_unified_sgd");
838
839        // Check that parameters were updated correctly
840        assert!((param.data[0] - 0.99).abs() < 1e-10);
841        assert!((param.data[1] - 1.98).abs() < 1e-10);
842        assert!((param.data[2] - 2.97).abs() < 1e-10);
843    }
844
845    #[test]
846    fn test_unified_adam() {
847        let config = OptimizerConfig::new(0.001f64);
848        let mut optimizer = UnifiedAdam::new(config);
849
850        let mut param = Parameter::new(Array1::from_vec(vec![1.0, 2.0, 3.0]), "test_param");
851        param.set_grad(Array1::from_vec(vec![0.1, 0.2, 0.3]));
852
853        optimizer
854            .step_param(&mut param)
855            .expect("optimizer.step_param succeeds in test_unified_adam");
856
857        // Parameters should have been updated (exact values depend on Adam's internal state)
858        assert!(param.data[0] < 1.0);
859        assert!(param.data[1] < 2.0);
860        assert!(param.data[2] < 3.0);
861    }
862
863    #[test]
864    fn test_optimizer_factory() {
865        let config = OptimizerConfig::new(0.01f64).weight_decay(0.0001);
866        let _sgd = OptimizerFactory::sgd(config.clone());
867        let _adam = OptimizerFactory::adam(config);
868    }
869
870    #[test]
871    fn test_parameter_operations() {
872        let mut param = Parameter::new(Array1::from_vec(vec![1.0, 2.0, 3.0]), "test");
873
874        // Test gradient setting
875        param.set_grad(Array1::from_vec(vec![0.1, 0.2, 0.3]));
876        assert!(param.grad().is_some());
877
878        // Test gradient clipping
879        param
880            .clip_grad(0.1)
881            .expect("param.clip_grad succeeds in test_parameter_operations");
882        let grad = param
883            .grad()
884            .expect("param.grad succeeds in test_parameter_operations");
885        let norm: f64 = grad.iter().map(|x| x * x).sum::<f64>().sqrt();
886        assert!((norm - 0.1).abs() < 1e-10);
887
888        // Test zero grad
889        param.zero_grad();
890        assert!(param.grad().is_none());
891    }
892
893    /// Regression test for the shared Adam step counter.
894    ///
895    /// `step_count` used to be a single counter incremented once per *parameter*, so
896    /// the second tensor updated in a training step was bias-corrected as if it were
897    /// on its second update. Every tensor must keep its own clock.
898    #[test]
899    fn test_unified_adam_step_count_is_per_parameter() {
900        let config = OptimizerConfig::new(0.1f64);
901        let mut optimizer = UnifiedAdam::new(config);
902
903        let mut first = Parameter::new(Array1::from_vec(vec![0.0f64]), "layer1.weight");
904        first.set_grad(Array1::from_vec(vec![1.0f64]));
905        let mut second = Parameter::new(Array1::from_vec(vec![0.0f64]), "layer2.weight");
906        second.set_grad(Array1::from_vec(vec![1.0f64]));
907
908        optimizer.step_param(&mut first).expect("first step failed");
909        optimizer
910            .step_param(&mut second)
911            .expect("second step failed");
912
913        assert_eq!(optimizer.step_count("layer1.weight"), 1);
914        assert_eq!(optimizer.step_count("layer2.weight"), 1);
915
916        // At t = 1 with a unit gradient the Adam step is -lr up to the epsilon term
917        // (denominator sqrt(v_hat) + eps), i.e. a relative error of about 3e-7.
918        assert!((first.data[0] + 0.1).abs() < 1e-6, "got {}", first.data[0]);
919        assert!(
920            (second.data[0] + 0.1).abs() < 1e-6,
921            "second tensor used the wrong timestep: {}",
922            second.data[0]
923        );
924        assert!((first.data[0] - second.data[0]).abs() < 1e-12);
925    }
926
927    /// A round trip through the state dictionary must reproduce the exact trajectory.
928    #[test]
929    fn test_unified_adam_state_dict_round_trip() {
930        let config = OptimizerConfig::new(0.05f64).weight_decay(0.01);
931        let mut original = UnifiedAdam::new(config.clone());
932
933        let mut param = Parameter::new(Array1::from_vec(vec![1.0f64, 2.0, 3.0]), "w");
934        for i in 0..5 {
935            let scale = 1.0 + i as f64;
936            param.set_grad(Array1::from_vec(vec![0.1 * scale, -0.2, 0.3]));
937            original.step_param(&mut param).expect("step failed");
938        }
939
940        let state = original.state_dict().expect("state_dict failed");
941        assert!(!state.is_empty(), "state dict must not be empty");
942        assert!(state.contains_key("adam.exp_avg.w"));
943        assert!(state.contains_key("adam.exp_avg_sq.w"));
944        assert!(state.contains_key("adam.step_count.w"));
945
946        let mut restored = UnifiedAdam::new(OptimizerConfig::new(999.0f64));
947        restored
948            .load_state_dict(state)
949            .expect("load_state_dict failed");
950
951        assert_eq!(restored.step_count("w"), 5);
952        assert!((restored.get_lr() - 0.05).abs() < 1e-12);
953
954        // Continue both optimizers from the same parameters and compare.
955        let mut a = param.clone();
956        let mut b = param.clone();
957        a.set_grad(Array1::from_vec(vec![0.4f64, -0.2, 0.3]));
958        b.set_grad(Array1::from_vec(vec![0.4f64, -0.2, 0.3]));
959        original.step_param(&mut a).expect("continue original");
960        restored.step_param(&mut b).expect("continue restored");
961
962        for i in 0..3 {
963            assert!(
964                (a.data[i] - b.data[i]).abs() < 1e-12,
965                "restored optimizer diverged at {}: {} vs {}",
966                i,
967                a.data[i],
968                b.data[i]
969            );
970        }
971    }
972
973    /// Loading a checkpoint whose buffers do not match must be rejected, not ignored.
974    #[test]
975    fn test_unified_adam_load_state_dict_validates_shapes() {
976        let mut optimizer = UnifiedAdam::new(OptimizerConfig::new(0.1f64));
977        let mut param = Parameter::new(Array1::from_vec(vec![1.0f64, 2.0, 3.0]), "w");
978        param.set_grad(Array1::from_vec(vec![0.1f64, 0.2, 0.3]));
979        optimizer.step_param(&mut param).expect("step failed");
980
981        let mut state = optimizer.state_dict().expect("state_dict failed");
982
983        // Shrink one moment buffer: the optimizer already holds three elements.
984        let mut truncated = state
985            .get("adam.exp_avg.w")
986            .cloned()
987            .expect("exp_avg entry must exist");
988        truncated.truncate(8);
989        state.insert("adam.exp_avg.w".to_string(), truncated);
990
991        assert!(optimizer.load_state_dict(state.clone()).is_err());
992
993        // A payload that is not a whole number of f64 values is rejected too.
994        state.insert("adam.exp_avg.w".to_string(), vec![0u8; 7]);
995        assert!(optimizer.load_state_dict(state).is_err());
996    }
997
998    /// SGD momentum buffers survive a state-dict round trip.
999    #[test]
1000    fn test_unified_sgd_state_dict_round_trip() {
1001        let config = OptimizerConfig::new(0.1f64);
1002        let mut original = UnifiedSGD::with_momentum(config, 0.9);
1003
1004        let mut param = Parameter::new(Array1::from_vec(vec![1.0f64, 2.0]), "w");
1005        for _ in 0..3 {
1006            param.set_grad(Array1::from_vec(vec![0.1f64, 0.2]));
1007            original.step_param(&mut param).expect("step failed");
1008        }
1009
1010        let state = original.state_dict().expect("state_dict failed");
1011        assert!(state.contains_key("sgd.momentum_buffer.w"));
1012
1013        let mut restored = UnifiedSGD::new(OptimizerConfig::new(999.0f64));
1014        restored
1015            .load_state_dict(state)
1016            .expect("load_state_dict failed");
1017        assert!((restored.get_lr() - 0.1).abs() < 1e-12);
1018
1019        let mut a = param.clone();
1020        let mut b = param.clone();
1021        a.set_grad(Array1::from_vec(vec![0.1f64, 0.2]));
1022        b.set_grad(Array1::from_vec(vec![0.1f64, 0.2]));
1023        original.step_param(&mut a).expect("continue original");
1024        restored.step_param(&mut b).expect("continue restored");
1025
1026        assert!((a.data[0] - b.data[0]).abs() < 1e-12);
1027        assert!((a.data[1] - b.data[1]).abs() < 1e-12);
1028    }
1029
1030    /// A state dict without a recognised version header must be rejected.
1031    #[test]
1032    fn test_state_dict_version_is_checked() {
1033        let mut optimizer = UnifiedSGD::new(OptimizerConfig::new(0.1f64));
1034        let mut state = HashMap::new();
1035        state.insert("config.lr".to_string(), 0.5f64.to_le_bytes().to_vec());
1036        assert!(optimizer.load_state_dict(state.clone()).is_err());
1037
1038        state.insert("format.version".to_string(), 7.0f64.to_le_bytes().to_vec());
1039        assert!(optimizer.load_state_dict(state).is_err());
1040    }
1041}
optirs_core/unified_api.rs

optirs_core/
unified_api.rs
