rustytorch/crates/training/rtx-transformers/optimizer_integration_test.rs

#!/usr/bin/env rust-script

//! Test for magnitude pruning optimizer integration
//! Tests the complete training workflow with gradual pruning and optimizer integration

use std::collections::HashMap;

/// Result type for this standalone test
type Result<T> = std::result::Result<T, Box<dyn std::error::Error>>;

/// Simple error type for testing
#[derive(Debug)]
struct TestError(String);

impl std::fmt::Display for TestError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "Test error: {}", self.0)
    }
}

impl std::error::Error for TestError {}

// Include all the implementations from our full magnitude pruning module

/// Configuration for magnitude-based pruning
#[derive(Debug, Clone, PartialEq)]
pub struct PruningConfig {
    /// Pruning strategy to apply
    pub strategy: PruningStrategy,
    /// Target sparsity ratio (0.0 to 1.0)
    pub sparsity_ratio: f64,
    /// Pruning granularity
    pub granularity: PruningGranularity,
    /// Optional pruning schedule for gradual pruning
    pub schedule: Option<PruningSchedule>,
    /// Random seed for deterministic pruning
    pub seed: Option<u64>,
}

/// Strategy for selecting weights to prune
#[derive(Debug, Clone, PartialEq)]
pub enum PruningStrategy {
    /// Global magnitude pruning across all parameters
    Global,
    /// Layer-wise magnitude pruning within each layer
    LayerWise,
    /// Structured pruning (channels, filters)
    Structured {
        /// Type of structured pattern
        pattern_type: StructuredPattern,
    },
}

/// Granularity of pruning operations
#[derive(Debug, Clone, PartialEq)]
pub enum PruningGranularity {
    /// Element-wise pruning (unstructured)
    Element,
    /// Channel-wise pruning
    Channel,
    /// Filter-wise pruning
    Filter,
    /// Block-wise pruning
    Block { height: usize, width: usize },
}

/// Structured pruning patterns
#[derive(Debug, Clone, PartialEq)]
pub enum StructuredPattern {
    /// Channel pruning
    Channel,
    /// Filter pruning
    Filter,
    /// N:M sparsity pattern
    NM { n: usize, m: usize },
}

/// Pruning schedule for gradual sparsification during training
#[derive(Debug, Clone, PartialEq)]
pub enum PruningSchedule {
    /// Linear schedule from initial to final sparsity
    Linear {
        /// Initial sparsity ratio
        initial_sparsity: f64,
        /// Final sparsity ratio  
        final_sparsity: f64,
        /// Duration in training steps
        duration_steps: usize,
    },
    /// Polynomial schedule
    Polynomial {
        /// Initial sparsity ratio
        initial_sparsity: f64,
        /// Final sparsity ratio
        final_sparsity: f64,
        /// Duration in training steps
        duration_steps: usize,
        /// Polynomial exponent
        exponent: f64,
    },
    /// Exponential schedule
    Exponential {
        /// Initial sparsity ratio
        initial_sparsity: f64,
        /// Final sparsity ratio
        final_sparsity: f64,
        /// Duration in training steps
        duration_steps: usize,
        /// Decay rate
        decay_rate: f64,
    },
}

/// Binary mask for indicating which weights to keep/prune
#[derive(Debug, Clone)]
pub struct PruningMask {
    /// Binary mask values (true = keep, false = prune)
    mask: Vec<bool>,
    /// Shape of the mask
    shape: Vec<usize>,
    /// Number of parameters kept
    kept_parameters: usize,
    /// Total number of parameters
    total_parameters: usize,
}

/// Statistics about pruning results
#[derive(Debug, Clone, PartialEq)]
pub struct PruningStats {
    /// Original parameter count
    pub original_params: usize,
    /// Parameters after pruning
    pub pruned_params: usize,
    /// Achieved sparsity ratio
    pub sparsity_ratio: f64,
    /// Compression ratio
    pub compression_ratio: f64,
    /// Memory savings (in bytes, estimated)
    pub memory_savings: usize,
}

/// Optimizer integration for pruned training
#[derive(Debug, Clone)]
pub struct PrunedOptimizer {
    /// Active pruning masks for each parameter group
    pub masks: HashMap<String, PruningMask>,
    /// Update count for scheduling
    pub update_count: usize,
    /// Pruning schedule configuration
    pub schedule_config: Option<PruningSchedule>,
}

/// Training state for pruned networks
#[derive(Debug, Clone)]
pub struct PrunedTrainingState {
    /// Current training step
    pub current_step: usize,
    /// Active sparsity ratio
    pub current_sparsity: f64,
    /// Accumulated pruning statistics
    pub stats: PruningStats,
    /// Whether gradual pruning is active
    pub gradual_pruning_active: bool,
}

/// Gradual pruning controller for training loops
#[derive(Debug)]
pub struct GradualPruningController {
    /// Pruning configuration
    config: PruningConfig,
    /// Current training state
    state: PrunedTrainingState,
    /// Parameter tensors being pruned
    parameters: HashMap<String, DenseTensor>,
    /// Current masks for each parameter
    masks: HashMap<String, PruningMask>,
}

/// Main struct for magnitude-based pruning operations
#[derive(Debug)]
pub struct MagnitudePruning;

/// Simple tensor for testing
#[derive(Debug, Clone)]
pub struct DenseTensor {
    pub data: Vec<f32>,
    pub shape: Vec<usize>,
}

impl DenseTensor {
    pub fn new(data: Vec<f32>, shape: Vec<usize>) -> Self {
        assert_eq!(data.len(), shape.iter().product::<usize>());
        Self { data, shape }
    }
    
    pub fn shape(&self) -> &[usize] {
        &self.shape
    }
    
    pub fn data(&self) -> &[f32] {
        &self.data
    }
    
    pub fn zeros(shape: Vec<usize>) -> Self {
        let total_elements = shape.iter().product::<usize>();
        Self {
            data: vec![0.0; total_elements],
            shape,
        }
    }
    
    /// Simulate parameter update (add gradients)
    pub fn update_with_gradients(&mut self, gradients: &DenseTensor, learning_rate: f32) -> Result<()> {
        if self.shape != gradients.shape {
            return Err(Box::new(TestError("Parameter and gradient shapes must match".into())));
        }
        
        for (param, grad) in self.data.iter_mut().zip(gradients.data.iter()) {
            *param -= learning_rate * grad;
        }
        
        Ok(())
    }
}

impl PruningConfig {
    /// Create global magnitude pruning configuration
    pub fn global(sparsity_ratio: f64) -> Self {
        Self {
            strategy: PruningStrategy::Global,
            sparsity_ratio,
            granularity: PruningGranularity::Element,
            schedule: None,
            seed: None,
        }
    }
    
    /// Create layer-wise magnitude pruning configuration
    pub fn layerwise(sparsity_ratio: f64) -> Self {
        Self {
            strategy: PruningStrategy::LayerWise,
            sparsity_ratio,
            granularity: PruningGranularity::Element,
            schedule: None,
            seed: None,
        }
    }
    
    /// Create structured magnitude pruning configuration
    pub fn structured(sparsity_ratio: f64, pattern: StructuredPattern) -> Self {
        let granularity = match pattern {
            StructuredPattern::Channel => PruningGranularity::Channel,
            StructuredPattern::Filter => PruningGranularity::Filter,
            StructuredPattern::NM { .. } => PruningGranularity::Element,
        };
        
        Self {
            strategy: PruningStrategy::Structured {
                pattern_type: pattern,
            },
            sparsity_ratio,
            granularity,
            schedule: None,
            seed: None,
        }
    }
    
    /// Add pruning schedule
    pub fn with_schedule(mut self, schedule: PruningSchedule) -> Self {
        self.schedule = Some(schedule);
        self
    }
    
    /// Validate configuration parameters
    pub fn validate(&self) -> Result<()> {
        if !(0.0..=1.0).contains(&self.sparsity_ratio) {
            return Err(Box::new(TestError("Sparsity ratio must be between 0.0 and 1.0".into())));
        }
        Ok(())
    }
}

impl PruningSchedule {
    /// Create linear pruning schedule
    pub fn linear(duration_steps: usize, final_sparsity: f64) -> Self {
        Self::Linear {
            initial_sparsity: 0.0,
            final_sparsity,
            duration_steps,
        }
    }
    
    /// Compute sparsity ratio at given training step
    pub fn compute_sparsity(&self, current_step: usize) -> f64 {
        match self {
            Self::Linear { initial_sparsity, final_sparsity, duration_steps } => {
                if current_step >= *duration_steps {
                    return *final_sparsity;
                }
                let progress = current_step as f64 / *duration_steps as f64;
                initial_sparsity + progress * (final_sparsity - initial_sparsity)
            }
            Self::Polynomial { initial_sparsity, final_sparsity, duration_steps, exponent } => {
                if current_step >= *duration_steps {
                    return *final_sparsity;
                }
                let progress = current_step as f64 / *duration_steps as f64;
                let poly_progress = progress.powf(*exponent);
                initial_sparsity + poly_progress * (final_sparsity - initial_sparsity)
            }
            Self::Exponential { initial_sparsity, final_sparsity, duration_steps, decay_rate } => {
                if current_step >= *duration_steps {
                    return *final_sparsity;
                }
                let progress = current_step as f64 / *duration_steps as f64;
                let exp_progress = 1.0 - (-decay_rate * progress).exp();
                initial_sparsity + exp_progress * (final_sparsity - initial_sparsity)
            }
        }
    }
}

impl PruningMask {
    /// Create new pruning mask
    pub fn new(mask: Vec<bool>, shape: Vec<usize>) -> Self {
        let total_parameters = shape.iter().product::<usize>();
        let kept_parameters = mask.iter().filter(|&&x| x).count();
        
        Self {
            mask,
            shape,
            kept_parameters,
            total_parameters,
        }
    }
    
    /// Get mask values
    pub fn mask(&self) -> &[bool] {
        &self.mask
    }
    
    /// Get mask shape
    pub fn shape(&self) -> &[usize] {
        &self.shape
    }
    
    /// Get sparsity ratio
    pub fn sparsity_ratio(&self) -> f64 {
        1.0 - (self.kept_parameters as f64 / self.total_parameters as f64)
    }
    
    /// Apply mask to tensor (element-wise multiplication)
    pub fn apply(&self, tensor: &DenseTensor) -> Result<DenseTensor> {
        if tensor.shape != self.shape {
            return Err(Box::new(TestError("Tensor and mask shapes must match".into())));
        }
        
        let masked_data: Vec<f32> = tensor.data.iter()
            .zip(self.mask.iter())
            .map(|(&value, &keep)| if keep { value } else { 0.0 })
            .collect();
        
        Ok(DenseTensor::new(masked_data, tensor.shape.clone()))
    }
    
    /// Apply mask to gradients (used in optimizer integration)
    pub fn apply_to_gradients(&self, gradients: &DenseTensor) -> Result<DenseTensor> {
        if gradients.shape != self.shape {
            return Err(Box::new(TestError("Gradient and mask shapes must match".into())));
        }
        
        // Zero out gradients for pruned weights
        let masked_gradients: Vec<f32> = gradients.data.iter()
            .zip(self.mask.iter())
            .map(|(&grad, &keep)| if keep { grad } else { 0.0 })
            .collect();
        
        Ok(DenseTensor::new(masked_gradients, gradients.shape.clone()))
    }
}

impl PruningStats {
    /// Create new pruning statistics
    pub fn new(original_params: usize, pruned_params: usize) -> Self {
        let sparsity_ratio = 1.0 - (pruned_params as f64 / original_params as f64);
        let compression_ratio = if pruned_params == 0 {
            f64::INFINITY
        } else {
            original_params as f64 / pruned_params as f64
        };
        
        // Estimate memory savings (assuming 4 bytes per f32)
        let memory_savings = (original_params - pruned_params) * 4;
        
        Self {
            original_params,
            pruned_params,
            sparsity_ratio,
            compression_ratio,
            memory_savings,
        }
    }
    
    /// Check if tensor meets target sparsity threshold
    pub fn meets_sparsity_target(&self, target_sparsity: f64) -> bool {
        self.sparsity_ratio >= target_sparsity
    }
    
    /// Merge statistics from multiple pruning operations
    pub fn merge(&self, other: &PruningStats) -> Self {
        let total_original = self.original_params + other.original_params;
        let total_pruned = self.pruned_params + other.pruned_params;
        Self::new(total_original, total_pruned)
    }
}

impl MagnitudePruning {
    /// Apply magnitude-based pruning to tensor
    pub fn prune(tensor: &DenseTensor, config: &PruningConfig) -> Result<DenseTensor> {
        config.validate()?;
        
        match config.strategy {
            PruningStrategy::Global => {
                Self::prune_global(tensor, config)
            }
            PruningStrategy::LayerWise => {
                Self::prune_layerwise(tensor, config)
            }
            PruningStrategy::Structured { ref pattern_type } => {
                Self::prune_structured(tensor, config, pattern_type)
            }
        }
    }
    
    /// Create pruning mask without modifying the original tensor
    pub fn create_mask(tensor: &DenseTensor, config: &PruningConfig) -> Result<PruningMask> {
        config.validate()?;
        
        match config.strategy {
            PruningStrategy::Global => {
                Self::create_global_mask(tensor, config)
            }
            PruningStrategy::LayerWise => {
                Self::create_layerwise_mask(tensor, config)
            }
            PruningStrategy::Structured { ref pattern_type } => {
                Self::create_structured_mask(tensor, config, pattern_type)
            }
        }
    }
    
    /// Compute importance scores based on magnitude
    pub fn compute_importance_scores(tensor: &DenseTensor) -> Result<Vec<f32>> {
        let scores: Vec<f32> = tensor.data.iter().map(|&x| x.abs()).collect();
        Ok(scores)
    }
    
    fn prune_global(tensor: &DenseTensor, config: &PruningConfig) -> Result<DenseTensor> {
        let mask = Self::create_global_mask(tensor, config)?;
        mask.apply(tensor)
    }
    
    fn prune_layerwise(tensor: &DenseTensor, config: &PruningConfig) -> Result<DenseTensor> {
        let mask = Self::create_layerwise_mask(tensor, config)?;
        mask.apply(tensor)
    }
    
    fn prune_structured(
        tensor: &DenseTensor,
        config: &PruningConfig,
        _pattern: &StructuredPattern,
    ) -> Result<DenseTensor> {
        // Simplified for test
        Self::prune_global(tensor, config)
    }
    
    fn create_global_mask(tensor: &DenseTensor, config: &PruningConfig) -> Result<PruningMask> {
        let total_elements = tensor.data.len();
        let elements_to_keep = ((1.0 - config.sparsity_ratio) * total_elements as f64).round() as usize;
        
        // Compute importance scores (absolute magnitudes)
        let importance_scores = Self::compute_importance_scores(tensor)?;
        
        // Create (index, score) pairs for sorting
        let mut indexed_scores: Vec<(usize, f32)> = importance_scores.iter()
            .enumerate()
            .map(|(i, &score)| (i, score))
            .collect();
        
        // Sort by importance score in descending order (keep highest magnitude values)
        indexed_scores.sort_by(|a, b| b.1.total_cmp(&a.1));
        
        // Create mask: true for kept elements, false for pruned elements
        let mut mask_data = vec![false; total_elements];
        for i in 0..elements_to_keep {
            if i < indexed_scores.len() {
                let (index, _) = indexed_scores[i];
                mask_data[index] = true;
            }
        }
        
        Ok(PruningMask::new(mask_data, tensor.shape.clone()))
    }
    
    fn create_layerwise_mask(tensor: &DenseTensor, config: &PruningConfig) -> Result<PruningMask> {
        // For layer-wise pruning, we currently support 2D tensors (treating each row as a layer)
        if tensor.shape.len() != 2 {
            return Err(Box::new(TestError("Layer-wise pruning currently only supports 2D tensors".into())));
        }
        
        let rows = tensor.shape[0];
        let cols = tensor.shape[1];
        let elements_per_row = cols;
        let elements_to_keep_per_row = ((1.0 - config.sparsity_ratio) * elements_per_row as f64).round() as usize;
        
        let mut mask_data = vec![false; tensor.data.len()];
        
        // Process each row independently
        for row in 0..rows {
            let row_start = row * cols;
            let row_end = row_start + cols;
            
            // Get magnitudes for this row
            let row_scores: Vec<(usize, f32)> = (row_start..row_end)
                .map(|i| (i, tensor.data[i].abs()))
                .collect();
            
            // Sort by magnitude in descending order
            let mut sorted_scores = row_scores;
            sorted_scores.sort_by(|a, b| b.1.total_cmp(&a.1));
            
            // Keep the top elements_to_keep_per_row elements in this row
            for i in 0..elements_to_keep_per_row.min(sorted_scores.len()) {
                let (original_index, _) = sorted_scores[i];
                mask_data[original_index] = true;
            }
        }
        
        Ok(PruningMask::new(mask_data, tensor.shape.clone()))
    }
    
    fn create_structured_mask(
        tensor: &DenseTensor,
        config: &PruningConfig,
        _pattern: &StructuredPattern,
    ) -> Result<PruningMask> {
        // Simplified for test - use global mask
        Self::create_global_mask(tensor, config)
    }
}

impl PrunedOptimizer {
    /// Create new pruned optimizer state
    pub fn new() -> Self {
        Self {
            masks: HashMap::new(),
            update_count: 0,
            schedule_config: None,
        }
    }
    
    /// Create pruned optimizer with schedule
    pub fn with_schedule(schedule: PruningSchedule) -> Self {
        Self {
            masks: HashMap::new(),
            update_count: 0,
            schedule_config: Some(schedule),
        }
    }
    
    /// Add parameter mask for tracking
    pub fn add_parameter_mask(&mut self, name: &str, mask: PruningMask) {
        self.masks.insert(name.to_string(), mask);
    }
    
    /// Get current sparsity ratio based on schedule
    pub fn get_current_sparsity(&self) -> f64 {
        if let Some(ref schedule) = self.schedule_config {
            schedule.compute_sparsity(self.update_count)
        } else {
            // If no schedule, return average sparsity of all masks
            if self.masks.is_empty() {
                0.0
            } else {
                let total_sparsity: f64 = self.masks.values()
                    .map(|mask| mask.sparsity_ratio())
                    .sum();
                total_sparsity / self.masks.len() as f64
            }
        }
    }
    
    /// Apply pruning masks to gradients before optimizer step
    pub fn apply_masks_to_gradients(&self, gradients: &HashMap<String, DenseTensor>) -> Result<HashMap<String, DenseTensor>> {
        let mut masked_gradients = HashMap::new();
        
        for (param_name, grad_tensor) in gradients {
            if let Some(mask) = self.masks.get(param_name) {
                let masked_grad = mask.apply_to_gradients(grad_tensor)?;
                masked_gradients.insert(param_name.clone(), masked_grad);
            } else {
                // No mask for this parameter, keep gradients as-is
                masked_gradients.insert(param_name.clone(), grad_tensor.clone());
            }
        }
        
        Ok(masked_gradients)
    }
    
    /// Update masks based on current parameters and schedule
    pub fn update_masks(&mut self, parameters: &HashMap<String, DenseTensor>, config: &PruningConfig) -> Result<()> {
        let target_sparsity = if let Some(ref schedule) = self.schedule_config {
            schedule.compute_sparsity(self.update_count)
        } else {
            config.sparsity_ratio
        };
        
        // Create updated config with current target sparsity
        let mut updated_config = config.clone();
        updated_config.sparsity_ratio = target_sparsity;
        
        // Update masks for each parameter
        for (param_name, tensor) in parameters {
            let new_mask = MagnitudePruning::create_mask(tensor, &updated_config)?;
            self.masks.insert(param_name.clone(), new_mask);
        }
        
        self.update_count += 1;
        Ok(())
    }
    
    /// Get pruning statistics across all parameters
    pub fn get_overall_stats(&self, parameters: &HashMap<String, DenseTensor>) -> Result<PruningStats> {
        let mut total_original = 0;
        let mut total_pruned = 0;
        
        for (param_name, tensor) in parameters {
            let original_count = tensor.data.len();
            total_original += original_count;
            
            if let Some(mask) = self.masks.get(param_name) {
                total_pruned += mask.kept_parameters;
            } else {
                // No mask means all parameters are kept
                total_pruned += original_count;
            }
        }
        
        Ok(PruningStats::new(total_original, total_pruned))
    }
}

impl PrunedTrainingState {
    /// Create new training state
    pub fn new() -> Self {
        Self {
            current_step: 0,
            current_sparsity: 0.0,
            stats: PruningStats::new(0, 0),
            gradual_pruning_active: false,
        }
    }
    
    /// Update training state
    pub fn update_step(&mut self, new_stats: PruningStats, sparsity: f64) {
        self.current_step += 1;
        self.current_sparsity = sparsity;
        self.stats = new_stats;
    }
    
    /// Check if gradual pruning should be applied
    pub fn should_update_pruning(&self, update_frequency: usize) -> bool {
        self.gradual_pruning_active && (self.current_step % update_frequency == 0)
    }
}

impl GradualPruningController {
    /// Create new gradual pruning controller
    pub fn new(config: PruningConfig) -> Self {
        Self {
            config,
            state: PrunedTrainingState::new(),
            parameters: HashMap::new(),
            masks: HashMap::new(),
        }
    }
    
    /// Add parameter to be pruned
    pub fn add_parameter(&mut self, name: &str, tensor: DenseTensor) {
        self.parameters.insert(name.to_string(), tensor);
    }
    
    /// Update parameters and recompute masks if needed
    pub fn update_parameters(&mut self, parameters: &HashMap<String, DenseTensor>) -> Result<()> {
        // Update stored parameters
        for (name, tensor) in parameters {
            self.parameters.insert(name.clone(), tensor.clone());
        }
        
        // Compute current sparsity based on schedule
        let current_sparsity = if let Some(ref schedule) = self.config.schedule {
            schedule.compute_sparsity(self.state.current_step)
        } else {
            self.config.sparsity_ratio
        };
        
        // Update config with current sparsity
        let mut current_config = self.config.clone();
        current_config.sparsity_ratio = current_sparsity;
        
        // Recompute masks for all parameters
        for (name, tensor) in &self.parameters {
            let mask = MagnitudePruning::create_mask(tensor, &current_config)?;
            self.masks.insert(name.clone(), mask);
        }
        
        // Update training state
        let overall_stats = self.compute_overall_stats();
        self.state.update_step(overall_stats, current_sparsity);
        
        Ok(())
    }
    
    /// Get current pruning masks
    pub fn get_masks(&self) -> &HashMap<String, PruningMask> {
        &self.masks
    }
    
    /// Get current training state
    pub fn get_state(&self) -> &PrunedTrainingState {
        &self.state
    }
    
    /// Apply current masks to parameters
    pub fn apply_pruning(&self, parameters: &HashMap<String, DenseTensor>) -> Result<HashMap<String, DenseTensor>> {
        let mut pruned_parameters = HashMap::new();
        
        for (name, tensor) in parameters {
            if let Some(mask) = self.masks.get(name) {
                let pruned_tensor = mask.apply(tensor)?;
                pruned_parameters.insert(name.clone(), pruned_tensor);
            } else {
                // No mask, keep original
                pruned_parameters.insert(name.clone(), tensor.clone());
            }
        }
        
        Ok(pruned_parameters)
    }
    
    /// Enable gradual pruning mode
    pub fn enable_gradual_pruning(&mut self) {
        self.state.gradual_pruning_active = true;
    }
    
    /// Disable gradual pruning mode
    pub fn disable_gradual_pruning(&mut self) {
        self.state.gradual_pruning_active = false;
    }
    
    /// Compute overall statistics
    fn compute_overall_stats(&self) -> PruningStats {
        let mut total_original = 0;
        let mut total_kept = 0;
        
        for (name, tensor) in &self.parameters {
            total_original += tensor.data.len();
            
            if let Some(mask) = self.masks.get(name) {
                total_kept += mask.kept_parameters;
            } else {
                total_kept += tensor.data.len();
            }
        }
        
        PruningStats::new(total_original, total_kept)
    }
}

/// Simulates a simple training loop with gradual pruning
fn simulate_training_with_pruning(
    initial_parameters: &HashMap<String, DenseTensor>,
    config: PruningConfig,
    num_steps: usize,
) -> Result<()> {
    println!("=== Starting Training Simulation ===");
    
    // Initialize gradual pruning controller
    let mut controller = GradualPruningController::new(config);
    controller.enable_gradual_pruning();
    
    // Add parameters
    for (name, tensor) in initial_parameters {
        controller.add_parameter(name, tensor.clone());
    }
    
    // Initialize optimizer with schedule
    let schedule = PruningSchedule::linear(num_steps, 0.8);
    let mut optimizer = PrunedOptimizer::with_schedule(schedule);
    
    // Current parameters (will be updated during training)
    let mut current_parameters = initial_parameters.clone();
    
    println!("Initial parameters:");
    for (name, tensor) in &current_parameters {
        println!("  {}: {} elements", name, tensor.data.len());
    }
    
    // Training loop
    for step in 0..num_steps {
        println!("\n--- Training Step {} ---", step);
        
        // Update controller with current parameters
        controller.update_parameters(&current_parameters)?;
        let state = controller.get_state();
        
        println!("Current sparsity: {:.2}%", state.current_sparsity * 100.0);
        println!("Active parameters: {}/{}", state.stats.pruned_params, state.stats.original_params);
        
        // Update optimizer masks
        let base_config = PruningConfig::global(0.0); // Will be overridden by schedule
        optimizer.update_masks(&current_parameters, &base_config)?;
        
        // Simulate gradients
        let mut gradients = HashMap::new();
        for (name, tensor) in &current_parameters {
            let grad_data: Vec<f32> = tensor.data.iter()
                .map(|_| (rand::random::<f32>() - 0.5) * 0.01) // Small random gradients
                .collect();
            gradients.insert(name.clone(), DenseTensor::new(grad_data, tensor.shape.clone()));
        }
        
        // Apply pruning masks to gradients
        let masked_gradients = optimizer.apply_masks_to_gradients(&gradients)?;
        
        // Update parameters with masked gradients
        let learning_rate = 0.01;
        for (name, param) in current_parameters.iter_mut() {
            if let Some(grad) = masked_gradients.get(name) {
                param.update_with_gradients(grad, learning_rate)?;
            }
        }
        
        // Apply pruning to ensure pruned weights remain zero
        current_parameters = controller.apply_pruning(&current_parameters)?;
        
        // Print statistics every 10 steps
        if step % 10 == 0 || step == num_steps - 1 {
            let stats = optimizer.get_overall_stats(&current_parameters)?;
            println!("Overall stats: {:.1}% sparsity, {:.1}x compression", 
                     stats.sparsity_ratio * 100.0, stats.compression_ratio);
        }
    }
    
    println!("\n=== Training Complete ===");
    let final_stats = optimizer.get_overall_stats(&current_parameters)?;
    println!("Final sparsity: {:.1}%", final_stats.sparsity_ratio * 100.0);
    println!("Final compression: {:.1}x", final_stats.compression_ratio);
    println!("Memory saved: {} bytes", final_stats.memory_savings);
    
    Ok(())
}

// Simple random number generator for testing
mod rand {
    use std::sync::atomic::{AtomicU32, Ordering};
    
    static SEED: AtomicU32 = AtomicU32::new(12345);
    
    pub fn random<T>() -> T
    where
        T: From<f32>,
    {
        let mut seed = SEED.load(Ordering::Relaxed);
        seed = seed.wrapping_mul(1103515245).wrapping_add(12345);
        SEED.store(seed, Ordering::Relaxed);
        let normalized = (seed as f32) / (u32::MAX as f32);
        T::from(normalized)
    }
}

fn main() -> Result<()> {
    println!("Running magnitude pruning optimizer integration tests...");
    
    // Test 1: Basic optimizer integration
    println!("\n=== Test 1: Basic Optimizer Integration ===");
    let mut optimizer = PrunedOptimizer::new();
    
    // Create test parameters
    let mut parameters = HashMap::new();
    parameters.insert("weight1".to_string(), DenseTensor::new(vec![0.1, 0.9, 0.2, 0.8], vec![4]));
    parameters.insert("weight2".to_string(), DenseTensor::new(vec![0.3, 0.7, 0.4, 0.6], vec![4]));
    
    // Create pruning masks
    let config = PruningConfig::global(0.5); // 50% sparsity
    for (name, tensor) in &parameters {
        let mask = MagnitudePruning::create_mask(tensor, &config)?;
        optimizer.add_parameter_mask(name, mask);
    }
    
    println!("Average sparsity: {:.1}%", optimizer.get_current_sparsity() * 100.0);
    
    // Create gradients
    let mut gradients = HashMap::new();
    gradients.insert("weight1".to_string(), DenseTensor::new(vec![0.01, 0.02, 0.03, 0.04], vec![4]));
    gradients.insert("weight2".to_string(), DenseTensor::new(vec![0.05, 0.06, 0.07, 0.08], vec![4]));
    
    // Apply masks to gradients
    let masked_gradients = optimizer.apply_masks_to_gradients(&gradients)?;
    println!("Original gradients for weight1: {:?}", gradients["weight1"].data);
    println!("Masked gradients for weight1: {:?}", masked_gradients["weight1"].data);
    
    let stats = optimizer.get_overall_stats(&parameters)?;
    println!("Overall stats: {:.1}% sparsity, {:.1}x compression", 
             stats.sparsity_ratio * 100.0, stats.compression_ratio);
    println!("✓ Basic optimizer integration passed");
    
    // Test 2: Gradual pruning with schedule
    println!("\n=== Test 2: Gradual Pruning with Schedule ===");
    let schedule = PruningSchedule::linear(50, 0.9); // Reach 90% sparsity in 50 steps
    let mut scheduled_optimizer = PrunedOptimizer::with_schedule(schedule);
    
    println!("Sparsity progression:");
    for step in [0, 10, 25, 50, 100] {
        let sparsity = scheduled_optimizer.get_current_sparsity();
        println!("  Step {}: {:.1}% sparsity", step, sparsity * 100.0);
        scheduled_optimizer.update_count = step;
    }
    println!("✓ Gradual pruning schedule passed");
    
    // Test 3: Training state management
    println!("\n=== Test 3: Training State Management ===");
    let mut state = PrunedTrainingState::new();
    state.gradual_pruning_active = true;
    
    // Simulate training steps
    for step in 1..=20 {
        let mock_stats = PruningStats::new(100, 100 - step * 2); // Gradual increase in sparsity
        state.update_step(mock_stats, step as f64 * 0.02);
        
        if step % 5 == 0 {
            println!("Step {}: {:.1}% sparsity, should update: {}", 
                     step, state.current_sparsity * 100.0, 
                     state.should_update_pruning(5));
        }
    }
    println!("✓ Training state management passed");
    
    // Test 4: Full training simulation
    println!("\n=== Test 4: Full Training Simulation ===");
    let mut simulation_params = HashMap::new();
    simulation_params.insert("layer1_weight".to_string(), 
                           DenseTensor::new((0..16).map(|i| (i as f32) * 0.1).collect(), vec![4, 4]));
    simulation_params.insert("layer2_weight".to_string(), 
                           DenseTensor::new((0..12).map(|i| (i as f32) * 0.05).collect(), vec![3, 4]));
    
    let sim_config = PruningConfig::global(0.0).with_schedule(PruningSchedule::linear(30, 0.8));
    simulate_training_with_pruning(&simulation_params, sim_config, 30)?;
    
    println!("✓ Full training simulation passed");
    
    // Test 5: Statistics merging
    println!("\n=== Test 5: Statistics Merging ===");
    let stats1 = PruningStats::new(100, 50);
    let stats2 = PruningStats::new(200, 80);
    let merged = stats1.merge(&stats2);
    
    println!("Stats1: {} -> {} ({:.1}% sparsity)", stats1.original_params, stats1.pruned_params, stats1.sparsity_ratio * 100.0);
    println!("Stats2: {} -> {} ({:.1}% sparsity)", stats2.original_params, stats2.pruned_params, stats2.sparsity_ratio * 100.0);
    println!("Merged: {} -> {} ({:.1}% sparsity)", merged.original_params, merged.pruned_params, merged.sparsity_ratio * 100.0);
    
    assert_eq!(merged.original_params, 300);
    assert_eq!(merged.pruned_params, 130);
    assert!(merged.meets_sparsity_target(0.5));
    println!("✓ Statistics merging passed");
    
    println!("\n🎉 All optimizer integration tests passed!");
    
    Ok(())
}