Struct TernaryOptimizerWrapper 

pub struct TernaryOptimizerWrapper { /* private fields */ }

Optimizer wrapper with ternary gradient accumulation.

Combines ternary accumulation with gradient management for memory-efficient training. Useful for large batch training where gradient accumulation is necessary.

Example

use vsa_optim_rs::ternary::TernaryOptimizerWrapper;
use vsa_optim_rs::TernaryConfig;

let mut wrapper = TernaryOptimizerWrapper::new(param_shapes, TernaryConfig::default(), &device)?;

for (i, batch) in batches.iter().enumerate() {
    // Compute gradients...
    let should_update = wrapper.step(&gradients)?;
    if should_update {
        let accumulated = wrapper.get_gradients_for_update()?;
        // Apply to optimizer...
    }
}

Implementations

impl TernaryOptimizerWrapper

pub fn new( param_shapes: &[(String, Vec<usize>)], config: TernaryConfig, device: &Device, ) -> Result<TernaryOptimizerWrapper, OptimError>

Create a new ternary optimizer wrapper.

Arguments

• param_shapes - List of (name, shape) tuples for parameters
• config - Ternary configuration
• device - Device for tensor storage

Errors

Returns error if accumulator creation fails.

pub fn step( &mut self, gradients: &HashMap<String, Tensor>, ) -> Result<bool, OptimError>

Accumulate gradient and check if update is needed.

Arguments

• gradients - Current step gradients

Returns

True if optimizer update should be performed, False if just accumulated.

Errors

Returns error if accumulation fails.

pub fn get_gradients_for_update( &mut self, ) -> Result<HashMap<String, Tensor>, OptimError>

Get accumulated gradients for optimizer update.

Call this when step() returns true.

Returns

Accumulated full-precision gradients.

Errors

Returns error if reconstruction fails.

pub fn get_stats(&self) -> OptimizerStats

Get optimization statistics.

pub const fn step_count(&self) -> usize

Get the step count.

pub const fn update_count(&self) -> usize

Get the update count.

pub fn reset_state(&mut self)

Reset state for checkpointing.

pub fn load_state(&mut self, step_count: usize, update_count: usize)

Load state from checkpoint.