Struct PhaseTrainer 

pub struct PhaseTrainer { /* private fields */ }

Orchestrates phase-based training for acceleration.

This is the main training orchestrator that combines all optimization techniques. It manages the phase transitions and ensures convergence while maximizing training speed.

The trainer automatically:

1. Tracks which phase we’re in
2. Manages gradient prediction during PREDICT phase
3. Applies corrections to prevent drift
4. Uses ternary accumulation for memory efficiency
5. Optionally uses VSA compression for gradient storage

Example

use vsa_optim_rs::phase::PhaseTrainer;
use vsa_optim_rs::PhaseConfig;

let shapes = vec![("layer.weight".to_string(), vec![64, 128])];
let config = PhaseConfig::default();
let mut trainer = PhaseTrainer::new(&shapes, config, &Device::Cpu)?;

// Training loop
for step in 0..total_steps {
    let step_info = trainer.begin_step()?;

    match step_info.phase {
        TrainingPhase::Full | TrainingPhase::Correct => {
            // Compute full gradients via backprop
            trainer.record_full_gradients(&gradients)?;
        }
        TrainingPhase::Predict => {
            // Use predicted gradients
            let predicted = trainer.get_predicted_gradients()?;
        }
    }

    trainer.end_step(loss_value)?;
}

Implementations

impl PhaseTrainer

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

Create a new phase trainer.

Arguments

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

Errors

Returns error if component initialization fails.

pub fn begin_step(&mut self) -> Result<StepInfo, OptimError>

Begin a training step. Returns info about current phase.

Returns

Step information including phase and whether phase changed.

Errors

Returns error if phase transition fails.

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

Record full gradients after backprop (for FULL or CORRECT phase).

Arguments

• gradients - Map of parameter names to gradient tensors

Errors

Returns error if recording fails.

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

Get predicted gradients (for PREDICT phase).

Returns

Map of parameter names to predicted gradient tensors.

Errors

Returns error if prediction fails.

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

Apply correction to gradients.

Arguments

• gradients - Mutable map of gradients to modify in-place

Errors

Returns error if correction fails.

pub fn end_step(&mut self, loss: f32) -> Result<(), OptimError>

End the training step.

Arguments

• loss - Loss value for this step

Errors

Returns error if tracking fails.

pub const fn current_phase(&self) -> TrainingPhase

Get current training phase.

pub const fn total_step(&self) -> usize

Get total step count.

pub const fn cycle_count(&self) -> usize

Get cycle count.

pub const fn speedup_ratio(&self) -> f32

Get speedup ratio.

pub fn get_stats(&self) -> TrainerStats

Get training statistics.

pub fn reset(&mut self) -> Result<(), OptimError>

Reset trainer state.

pub fn vsa_compressor_mut(&mut self) -> &mut VSAGradientCompressor

Get mutable access to VSA compressor.

pub fn ternary_accumulator_mut(&mut self) -> &mut TernaryGradientAccumulator

Get mutable access to ternary accumulator.

pub fn should_compute_full(&self) -> bool

Check if should compute full gradients.