foxhunt/ml/tests/tft_lstm_int8_quantization_test.rs

//! TFT LSTM Encoder INT8 Quantization Tests (TDD)
//!
//! Test suite for INT8 quantization of TFT's LSTM encoder layer.
//! Target: 800MB → 200MB (75% reduction) with <5% accuracy loss.
//!
//! ## LSTM Architecture
//! - 2 layers with hidden_dim=128
//! - 8 weight matrices per layer: Wii, Wif, Wig, Wio, Whi, Whf, Whg, Who
//! - Recurrent connections require careful quantization
//!
//! ## Test Strategy
//! 1. Quantize all 8 weight matrices per layer (16 total)
//! 2. Verify forward pass maintains temporal coherence
//! 3. Test hidden state shapes preserved
//! 4. Validate accuracy loss <5% on sequence prediction
//! 5. Confirm memory reduction 70-80%

use anyhow::Result;
use candle_core::{Device, Tensor};
use ml::memory_optimization::quantization::{QuantizationConfig, QuantizationType, Quantizer};
use ml::tft::lstm_encoder::LSTMEncoder;
use ml::tft::quantized_lstm::QuantizedLSTMEncoder;

#[test]
fn test_lstm_encoder_exists() -> Result<()> {
    // Test that LSTMEncoder struct exists (will fail initially)
    let device = Device::Cpu;
    let num_layers = 2;
    let input_size = 64;
    let hidden_size = 128;

    let encoder = LSTMEncoder::new(num_layers, input_size, hidden_size, &device)?;

    assert_eq!(encoder.num_layers(), num_layers);
    assert_eq!(encoder.hidden_size(), hidden_size);

    Ok(())
}

#[test]
fn test_quantized_lstm_encoder_creation() -> Result<()> {
    // Test creating quantized LSTM from FP32 model
    let device = Device::Cpu;
    let lstm = LSTMEncoder::new(2, 64, 128, &device)?;

    let config = QuantizationConfig {
        quant_type: QuantizationType::Int8,
        symmetric: true,
        per_channel: true,
        calibration_samples: Some(1000),
    };

    let quantized = QuantizedLSTMEncoder::from_f32_model(&lstm, config)?;

    assert_eq!(quantized.num_layers(), 2);
    assert_eq!(quantized.hidden_size(), 128);

    Ok(())
}

#[test]
fn test_quantize_all_lstm_weights() -> Result<()> {
    // Test that all 8 weight matrices per layer are quantized
    let device = Device::Cpu;
    let lstm = LSTMEncoder::new(2, 64, 128, &device)?;

    let config = QuantizationConfig {
        quant_type: QuantizationType::Int8,
        symmetric: true,
        per_channel: true,
        calibration_samples: Some(1000),
    };

    let quantized = QuantizedLSTMEncoder::from_f32_model(&lstm, config)?;

    // Verify all weight tensors are quantized
    let quantized_weights = quantized.get_quantized_weights();

    // 2 layers × 8 weight matrices = 16 total
    assert_eq!(quantized_weights.len(), 2, "Should have 2 layers");

    for (layer_idx, layer_weights) in quantized_weights.iter().enumerate() {
        // Each layer has 8 weight matrices
        assert_eq!(
            layer_weights.len(),
            8,
            "Layer {} should have 8 weight matrices (Wii, Wif, Wig, Wio, Whi, Whf, Whg, Who)",
            layer_idx
        );

        // Verify each weight is quantized to int8
        for (weight_name, quantized_tensor) in layer_weights.iter() {
            assert_eq!(
                quantized_tensor.quant_type,
                QuantizationType::Int8,
                "Weight {} in layer {} should be INT8",
                weight_name,
                layer_idx
            );
        }
    }

    Ok(())
}

#[test]
fn test_quantized_lstm_forward_pass() -> Result<()> {
    // Test forward pass maintains temporal coherence
    let device = Device::Cpu;
    let batch_size = 4;
    let seq_len = 20;
    let input_size = 64;
    let hidden_size = 128;

    // Create FP32 LSTM
    let lstm = LSTMEncoder::new(2, input_size, hidden_size, &device)?;

    // Create quantized version
    let config = QuantizationConfig {
        quant_type: QuantizationType::Int8,
        symmetric: true,
        per_channel: true,
        calibration_samples: Some(1000),
    };
    let quantized = QuantizedLSTMEncoder::from_f32_model(&lstm, config)?;

    // Create input: [batch, seq_len, input_size]
    let input = Tensor::randn(0f32, 1.0, (batch_size, seq_len, input_size), &device)?;

    // Forward pass through quantized LSTM
    let (output, hidden_state, cell_state) = quantized.forward(&input, None)?;

    // Verify output shape: [batch, seq_len, hidden_size]
    assert_eq!(output.dims(), &[batch_size, seq_len, hidden_size]);

    // Verify hidden state shape: [num_layers, batch, hidden_size]
    assert_eq!(hidden_state.dims(), &[2, batch_size, hidden_size]);

    // Verify cell state shape: [num_layers, batch, hidden_size]
    assert_eq!(cell_state.dims(), &[2, batch_size, hidden_size]);

    // Verify temporal coherence (no NaNs or Infs)
    let output_vec = output.flatten_all()?.to_vec1::<f32>()?;
    assert!(
        output_vec.iter().all(|x| x.is_finite()),
        "Output contains NaN or Inf"
    );

    Ok(())
}

#[test]
fn test_hidden_state_shapes_preserved() -> Result<()> {
    // Test that hidden state dimensions match original LSTM
    let device = Device::Cpu;
    let batch_size = 8;
    let seq_len = 15;
    let input_size = 64;
    let hidden_size = 128;
    let num_layers = 2;

    // Create both FP32 and quantized LSTMs
    let lstm_f32 = LSTMEncoder::new(num_layers, input_size, hidden_size, &device)?;
    let config = QuantizationConfig::default();
    let lstm_int8 = QuantizedLSTMEncoder::from_f32_model(&lstm_f32, config)?;

    // Create input
    let input = Tensor::randn(0f32, 1.0, (batch_size, seq_len, input_size), &device)?;

    // Forward pass through both
    let (output_f32, h_f32, c_f32) = lstm_f32.forward(&input, None)?;
    let (output_int8, h_int8, c_int8) = lstm_int8.forward(&input, None)?;

    // Verify shapes match exactly
    assert_eq!(
        output_f32.dims(),
        output_int8.dims(),
        "Output shapes must match"
    );
    assert_eq!(
        h_f32.dims(),
        h_int8.dims(),
        "Hidden state shapes must match"
    );
    assert_eq!(c_f32.dims(), c_int8.dims(), "Cell state shapes must match");

    Ok(())
}

#[test]
fn test_quantization_accuracy_loss_within_5_percent() -> Result<()> {
    // Test that quantization introduces <5% accuracy loss on sequence prediction
    let device = Device::Cpu;
    let batch_size = 16;
    let seq_len = 30;
    let input_size = 64;
    let hidden_size = 128;

    // Create FP32 LSTM
    let lstm_f32 = LSTMEncoder::new(2, input_size, hidden_size, &device)?;

    // Create quantized LSTM
    let config = QuantizationConfig {
        quant_type: QuantizationType::Int8,
        symmetric: true,
        per_channel: true,
        calibration_samples: Some(1000),
    };
    let lstm_int8 = QuantizedLSTMEncoder::from_f32_model(&lstm_f32, config)?;

    // Generate test sequences
    let num_samples = 100;
    let mut total_mse_f32 = 0.0;
    let mut total_mse_int8 = 0.0;

    for _ in 0..num_samples {
        // Create input sequence
        let input = Tensor::randn(0f32, 1.0, (batch_size, seq_len, input_size), &device)?;

        // Create synthetic target (next token prediction)
        let target = Tensor::randn(0f32, 1.0, (batch_size, seq_len, hidden_size), &device)?;

        // Forward pass through both LSTMs
        let (output_f32, _, _) = lstm_f32.forward(&input, None)?;
        let (output_int8, _, _) = lstm_int8.forward(&input, None)?;

        // Compute MSE for both
        let diff_f32 = (output_f32.clone() - &target)?;
        let mse_f32 = diff_f32.sqr()?.mean_all()?.to_scalar::<f32>()?;

        let diff_int8 = (output_int8 - &target)?;
        let mse_int8 = diff_int8.sqr()?.mean_all()?.to_scalar::<f32>()?;

        total_mse_f32 += mse_f32 as f64;
        total_mse_int8 += mse_int8 as f64;
    }

    let avg_mse_f32 = total_mse_f32 / num_samples as f64;
    let avg_mse_int8 = total_mse_int8 / num_samples as f64;

    // Calculate accuracy degradation
    let accuracy_loss = ((avg_mse_int8 - avg_mse_f32) / avg_mse_f32).abs() * 100.0;

    println!("FP32 MSE: {:.6}", avg_mse_f32);
    println!("INT8 MSE: {:.6}", avg_mse_int8);
    println!("Accuracy loss: {:.2}%", accuracy_loss);

    assert!(
        accuracy_loss < 5.0,
        "Accuracy loss {:.2}% exceeds 5% threshold",
        accuracy_loss
    );

    Ok(())
}

#[test]
fn test_memory_reduction_70_to_80_percent() -> Result<()> {
    // Test memory reduction is between 70-80%
    let device = Device::Cpu;
    let lstm_f32 = LSTMEncoder::new(2, 64, 128, &device)?;

    // Calculate FP32 memory usage
    let memory_f32_mb = lstm_f32.estimate_memory_mb();

    // Create quantized version
    let config = QuantizationConfig::default();
    let lstm_int8 = QuantizedLSTMEncoder::from_f32_model(&lstm_f32, config)?;

    // Calculate INT8 memory usage
    let memory_int8_mb = lstm_int8.estimate_memory_mb();

    // Calculate reduction percentage
    let reduction_percent = ((memory_f32_mb - memory_int8_mb) / memory_f32_mb) * 100.0;

    println!("FP32 memory: {:.2} MB", memory_f32_mb);
    println!("INT8 memory: {:.2} MB", memory_int8_mb);
    println!("Memory reduction: {:.2}%", reduction_percent);

    assert!(
        reduction_percent >= 70.0 && reduction_percent <= 80.0,
        "Memory reduction {:.2}% is outside 70-80% target range",
        reduction_percent
    );

    Ok(())
}

#[test]
fn test_quantized_lstm_with_initial_hidden_state() -> Result<()> {
    // Test forward pass with provided initial hidden/cell states
    let device = Device::Cpu;
    let batch_size = 4;
    let seq_len = 10;
    let input_size = 64;
    let hidden_size = 128;
    let num_layers = 2;

    let lstm = LSTMEncoder::new(num_layers, input_size, hidden_size, &device)?;
    let config = QuantizationConfig::default();
    let quantized = QuantizedLSTMEncoder::from_f32_model(&lstm, config)?;

    // Create input
    let input = Tensor::randn(0f32, 1.0, (batch_size, seq_len, input_size), &device)?;

    // Create initial hidden/cell states
    let h0 = Tensor::randn(0f32, 1.0, (num_layers, batch_size, hidden_size), &device)?;
    let c0 = Tensor::randn(0f32, 1.0, (num_layers, batch_size, hidden_size), &device)?;

    // Forward pass with initial states
    let (output, h_final, c_final) = quantized.forward(&input, Some((h0.clone(), c0.clone())))?;

    // Verify shapes
    assert_eq!(output.dims(), &[batch_size, seq_len, hidden_size]);
    assert_eq!(h_final.dims(), &[num_layers, batch_size, hidden_size]);
    assert_eq!(c_final.dims(), &[num_layers, batch_size, hidden_size]);

    // Verify states are different from initial (LSTM updated them)
    let h_diff = (h_final - &h0)?;
    let h_diff_norm = h_diff.sqr()?.sum_all()?.to_scalar::<f32>()?;
    assert!(h_diff_norm > 0.0, "Hidden state should be updated");

    Ok(())
}

#[test]
fn test_batch_size_independence() -> Result<()> {
    // Test that batch size doesn't affect per-sample output
    let device = Device::Cpu;
    let seq_len = 15;
    let input_size = 64;
    let hidden_size = 128;

    let lstm = LSTMEncoder::new(2, input_size, hidden_size, &device)?;
    let config = QuantizationConfig::default();
    let quantized = QuantizedLSTMEncoder::from_f32_model(&lstm, config)?;

    // Create single sample input
    let input_single = Tensor::randn(0f32, 1.0, (1, seq_len, input_size), &device)?;
    let (output_single, _, _) = quantized.forward(&input_single, None)?;

    // Create batched input (repeat the same sample 4 times)
    let input_batched = input_single.repeat(&[4, 1, 1])?;
    let (output_batched, _, _) = quantized.forward(&input_batched, None)?;

    // Extract first sample from batched output
    let output_batched_first = output_batched.narrow(0, 0, 1)?;

    // Compare with single-sample output
    let diff = (output_single - output_batched_first)?;
    let max_diff = diff.abs()?.max_all()?.to_scalar::<f32>()?;

    assert!(
        max_diff < 1e-5,
        "Batch size should not affect per-sample output (max diff: {})",
        max_diff
    );

    Ok(())
}

#[test]
fn test_quantization_config_options() -> Result<()> {
    // Test different quantization configurations
    let device = Device::Cpu;
    let lstm = LSTMEncoder::new(2, 64, 128, &device)?;

    // Symmetric quantization
    let config_symmetric = QuantizationConfig {
        quant_type: QuantizationType::Int8,
        symmetric: true,
        per_channel: true,
        calibration_samples: Some(1000),
    };
    let quantized_symmetric = QuantizedLSTMEncoder::from_f32_model(&lstm, config_symmetric)?;
    assert_eq!(quantized_symmetric.num_layers(), 2);

    // Asymmetric quantization
    let config_asymmetric = QuantizationConfig {
        quant_type: QuantizationType::Int8,
        symmetric: false,
        per_channel: true,
        calibration_samples: Some(1000),
    };
    let quantized_asymmetric = QuantizedLSTMEncoder::from_f32_model(&lstm, config_asymmetric)?;
    assert_eq!(quantized_asymmetric.num_layers(), 2);

    // Per-tensor quantization
    let config_per_tensor = QuantizationConfig {
        quant_type: QuantizationType::Int8,
        symmetric: true,
        per_channel: false,
        calibration_samples: Some(1000),
    };
    let quantized_per_tensor = QuantizedLSTMEncoder::from_f32_model(&lstm, config_per_tensor)?;
    assert_eq!(quantized_per_tensor.num_layers(), 2);

    Ok(())
}