foxhunt/ml/tests/tft_attention_gradient_flow.rs

//! # Wave 8.7: TFT Attention Mechanism Gradient Flow Tests
//!
//! Comprehensive test suite validating gradient flow through TFT attention mechanism.
//! These tests verify that gradients propagate correctly through:
//! 1. Multi-head attention (all heads)
//! 2. Q/K/V projection layers
//! 3. Positional encoding
//! 4. Causal masking
//! 5. Layer normalization and residual connections
//!
//! Context: Wave 7.3 verified no `.detach()` calls blocking gradients.
//! This wave adds comprehensive gradient flow validation.

#![allow(unused_crate_dependencies)]

use candle_core::{DType, Device, Tensor, Var};
use candle_nn::{Linear, Optimizer, VarBuilder, VarMap};

use ml::tft::temporal_attention::{AttentionHead, PositionalEncoding, TemporalSelfAttention};
use ml::MLError;

// ============================================================================
// HELPER FUNCTIONS FOR GRADIENT VERIFICATION
// ============================================================================

/// Helper function to check if gradients exist and are valid
fn verify_gradients(
    var: &Var,
    expected_min_norm: f32,
    test_name: &str,
) -> Result<(), MLError> {
    let grad = var
        .grad()
        .ok_or_else(|| MLError::ModelError(format!("{}: No gradient found", test_name)))?;

    let grad_vec = grad.flatten_all()?.to_vec1::<f32>()?;
    let grad_norm: f32 = grad_vec.iter().map(|&g| g.powi(2)).sum::<f32>().sqrt();

    assert!(
        grad_norm > expected_min_norm,
        "{}: Gradient norm {:.6} should be > {:.6}",
        test_name,
        grad_norm,
        expected_min_norm
    );
    assert!(
        !grad_norm.is_nan(),
        "{}: Gradient should not be NaN",
        test_name
    );
    assert!(
        !grad_norm.is_infinite(),
        "{}: Gradient should not be Inf",
        test_name
    );

    Ok(())
}

/// Helper to extract gradient norm from a tensor
fn compute_gradient_norm(tensor: &Tensor) -> Result<f32, MLError> {
    let grad_vec = tensor.flatten_all()?.to_vec1::<f32>()?;
    let grad_norm: f32 = grad_vec.iter().map(|&g| g.powi(2)).sum::<f32>().sqrt();
    Ok(grad_norm)
}

// ============================================================================
// TEST 1: BASIC ATTENTION INPUT GRADIENT FLOW
// ============================================================================

#[test]
fn test_attention_input_gradient_flow() -> Result<(), MLError> {
    let device = Device::Cpu;
    let varmap = VarMap::new();
    let vs = VarBuilder::from_varmap(&varmap, DType::F32, &device);

    // Create attention module
    let attention = TemporalSelfAttention::new(
        64,   // hidden_dim
        4,    // num_heads
        0.1,  // dropout_rate
        false, // use_flash_attention (disable for gradient testing)
        vs,
    )?;

    // Create input tensor as Var for gradient tracking
    let input_data = vec![0.1f32; 640]; // 2 * 10 * 32
    let input = Var::from_slice(&input_data, (2, 10, 64), &device)?;

    // Forward pass with causal masking
    let output = attention.forward(&input, true)?;

    // Compute loss (sum for testing)
    let loss = output.sum_all()?;

    // Backward pass
    let grads = loss.backward()?;

    // Verify gradients exist for input
    let input_grad = grads
        .get(&input)
        .ok_or_else(|| MLError::ModelError("No gradient for input".to_string()))?;

    // Check gradient is non-zero and valid
    let grad_norm = compute_gradient_norm(input_grad)?;

    println!("Test 1 - Input gradient norm: {:.6}", grad_norm);
    assert!(
        grad_norm > 0.001,
        "Input gradient norm should be > 0.001, got {:.6}",
        grad_norm
    );
    assert!(!grad_norm.is_nan(), "Input gradient should not be NaN");
    assert!(
        !grad_norm.is_infinite(),
        "Input gradient should not be Inf"
    );

    Ok(())
}

// ============================================================================
// TEST 2: MULTI-HEAD GRADIENT FLOW
// ============================================================================

#[test]
fn test_multihead_attention_gradient_flow() -> Result<(), MLError> {
    let device = Device::Cpu;
    let varmap = VarMap::new();
    let vs = VarBuilder::from_varmap(&varmap, DType::F32, &device);

    // Create attention with 8 heads
    let attention = TemporalSelfAttention::new(128, 8, 0.1, false, vs)?;

    // Create input
    let input_data = vec![0.2f32; 2560]; // 2 * 10 * 128
    let input = Var::from_slice(&input_data, (2, 10, 128), &device)?;

    // Forward pass
    let output = attention.forward(&input, false)?; // No causal mask

    // Compute loss
    let loss = output.sum_all()?;

    // Backward pass
    let grads = loss.backward()?;

    // Verify input gradients
    let input_grad = grads
        .get(&input)
        .ok_or_else(|| MLError::ModelError("No gradient for input".to_string()))?;

    let grad_norm = compute_gradient_norm(input_grad)?;
    println!("Test 2 - Multi-head gradient norm: {:.6}", grad_norm);

    assert!(
        grad_norm > 0.001,
        "Multi-head gradient should be non-zero: {:.6}",
        grad_norm
    );

    // Verify all attention head parameters have gradients
    let all_vars = varmap.all_vars();
    println!("Test 2 - Total trainable variables: {}", all_vars.len());

    // Check that projection layers (query, key, value) have gradients
    let mut projection_grads_found = 0;
    for var in all_vars {
        if let Some(grad) = var.grad() {
            let grad_norm = compute_gradient_norm(&grad)?;
            if grad_norm > 1e-6 {
                projection_grads_found += 1;
            }
        }
    }

    println!(
        "Test 2 - Projection layers with gradients: {}",
        projection_grads_found
    );
    assert!(
        projection_grads_found > 0,
        "At least some projection layers should have gradients"
    );

    Ok(())
}

// ============================================================================
// TEST 3: Q/K/V PROJECTION GRADIENT FLOW
// ============================================================================

#[test]
fn test_qkv_projection_gradient_flow() -> Result<(), MLError> {
    let device = Device::Cpu;

    // Create a single attention head directly for testing
    let head = AttentionHead::new(
        64,  // hidden_dim
        16,  // head_dim
        &device,
    )?;

    // Create input
    let input_data = vec![0.1f32; 320]; // 2 * 10 * 16
    let input = Var::from_slice(&input_data, (2, 10, 64), &device)?;

    // Forward pass
    let (output, _attention_weights) = head.forward(&input, None, 1.0)?;

    // Compute loss
    let loss = output.sum_all()?;

    // Backward pass
    let grads = loss.backward()?;

    // Verify input gradient
    let input_grad = grads
        .get(&input)
        .ok_or_else(|| MLError::ModelError("No gradient for input".to_string()))?;

    let grad_norm = compute_gradient_norm(input_grad)?;
    println!("Test 3 - Q/K/V projection gradient norm: {:.6}", grad_norm);

    assert!(
        grad_norm > 0.001,
        "Q/K/V projection gradient should be non-zero: {:.6}",
        grad_norm
    );

    Ok(())
}

// ============================================================================
// TEST 4: CAUSAL MASKING GRADIENT FLOW
// ============================================================================

#[test]
fn test_causal_masking_gradient_flow() -> Result<(), MLError> {
    let device = Device::Cpu;
    let varmap = VarMap::new();
    let vs = VarBuilder::from_varmap(&varmap, DType::F32, &device);

    // Create attention
    let attention = TemporalSelfAttention::new(64, 4, 0.1, false, vs)?;

    // Create input
    let input_data = vec![0.15f32; 640]; // 2 * 10 * 64
    let input = Var::from_slice(&input_data, (2, 10, 64), &device)?;

    // Forward pass WITH causal mask
    let output_causal = attention.forward(&input, true)?;
    let loss_causal = output_causal.sum_all()?;
    let grads_causal = loss_causal.backward()?;

    // Get gradient norm with causal masking
    let grad_causal = grads_causal
        .get(&input)
        .ok_or_else(|| MLError::ModelError("No gradient with causal mask".to_string()))?;
    let grad_norm_causal = compute_gradient_norm(grad_causal)?;

    println!("Test 4 - Causal masking gradient norm: {:.6}", grad_norm_causal);

    // Verify gradients flow with causal masking
    assert!(
        grad_norm_causal > 0.001,
        "Causal masking should preserve gradient flow: {:.6}",
        grad_norm_causal
    );
    assert!(
        !grad_norm_causal.is_nan(),
        "Gradient with causal mask should not be NaN"
    );

    Ok(())
}

// ============================================================================
// TEST 5: POSITIONAL ENCODING GRADIENT FLOW
// ============================================================================

#[test]
fn test_positional_encoding_gradient_flow() -> Result<(), MLError> {
    let device = Device::Cpu;
    let varmap = VarMap::new();
    let vs = VarBuilder::from_varmap(&varmap, DType::F32, &device);

    // Create attention with positional encoding
    let attention = TemporalSelfAttention::new(64, 4, 0.1, false, vs)?;

    // Create input
    let input_data = vec![0.1f32; 1280]; // 2 * 10 * 64
    let input = Var::from_slice(&input_data, (2, 10, 64), &device)?;

    // Forward pass (positional encoding is added inside forward())
    let output = attention.forward(&input, false)?;

    // Compute loss
    let loss = output.sum_all()?;

    // Backward pass
    let grads = loss.backward()?;

    // Verify input gradient (should receive gradients through positional encoding addition)
    let input_grad = grads
        .get(&input)
        .ok_or_else(|| MLError::ModelError("No gradient for input".to_string()))?;

    let grad_norm = compute_gradient_norm(input_grad)?;
    println!(
        "Test 5 - Positional encoding gradient norm: {:.6}",
        grad_norm
    );

    assert!(
        grad_norm > 0.001,
        "Positional encoding should not block gradients: {:.6}",
        grad_norm
    );

    Ok(())
}

// ============================================================================
// TEST 6: RESIDUAL CONNECTION GRADIENT FLOW
// ============================================================================

#[test]
fn test_residual_connection_gradient_flow() -> Result<(), MLError> {
    let device = Device::Cpu;
    let varmap = VarMap::new();
    let vs = VarBuilder::from_varmap(&varmap, DType::F32, &device);

    // Create attention
    let attention = TemporalSelfAttention::new(64, 4, 0.0, false, vs)?; // No dropout for testing

    // Create input with distinct values
    let mut input_data = Vec::new();
    for i in 0..640 {
        input_data.push((i as f32) * 0.01);
    }
    let input = Var::from_slice(&input_data, (2, 10, 64), &device)?;

    // Forward pass (includes residual connection: output = LayerNorm(input + attention(input)))
    let output = attention.forward(&input, false)?;

    // Compute loss
    let loss = output.sum_all()?;

    // Backward pass
    let grads = loss.backward()?;

    // Verify input gradient exists (should receive gradients through residual connection)
    let input_grad = grads
        .get(&input)
        .ok_or_else(|| MLError::ModelError("No gradient for input".to_string()))?;

    let grad_norm = compute_gradient_norm(input_grad)?;
    println!(
        "Test 6 - Residual connection gradient norm: {:.6}",
        grad_norm
    );

    assert!(
        grad_norm > 0.001,
        "Residual connection should preserve gradients: {:.6}",
        grad_norm
    );

    Ok(())
}

// ============================================================================
// TEST 7: LAYER NORMALIZATION GRADIENT FLOW
// ============================================================================

#[test]
fn test_layer_normalization_gradient_flow() -> Result<(), MLError> {
    let device = Device::Cpu;
    let varmap = VarMap::new();
    let vs = VarBuilder::from_varmap(&varmap, DType::F32, &device);

    // Create attention (LayerNorm is applied at the end)
    let attention = TemporalSelfAttention::new(64, 4, 0.1, false, vs)?;

    // Create input
    let input_data = vec![0.2f32; 640]; // 2 * 10 * 64
    let input = Var::from_slice(&input_data, (2, 10, 64), &device)?;

    // Forward pass
    let output = attention.forward(&input, false)?;

    // Compute loss
    let loss = output.sum_all()?;

    // Backward pass
    let grads = loss.backward()?;

    // Verify input gradient
    let input_grad = grads
        .get(&input)
        .ok_or_else(|| MLError::ModelError("No gradient for input".to_string()))?;

    let grad_norm = compute_gradient_norm(input_grad)?;
    println!(
        "Test 7 - Layer normalization gradient norm: {:.6}",
        grad_norm
    );

    assert!(
        grad_norm > 0.001,
        "Layer normalization should not block gradients: {:.6}",
        grad_norm
    );

    // Check LayerNorm parameters have gradients
    let all_vars = varmap.all_vars();
    let mut layernorm_grads = 0;

    for var in all_vars {
        if let Some(grad) = var.grad() {
            let grad_norm = compute_gradient_norm(&grad)?;
            if grad_norm > 1e-6 {
                layernorm_grads += 1;
            }
        }
    }

    println!(
        "Test 7 - LayerNorm parameters with gradients: {}",
        layernorm_grads
    );
    assert!(
        layernorm_grads > 0,
        "LayerNorm parameters should have gradients"
    );

    Ok(())
}

// ============================================================================
// TEST 8: DROPOUT GRADIENT FLOW
// ============================================================================

#[test]
fn test_dropout_gradient_flow() -> Result<(), MLError> {
    let device = Device::Cpu;
    let varmap = VarMap::new();
    let vs = VarBuilder::from_varmap(&varmap, DType::F32, &device);

    // Create attention with dropout
    let attention = TemporalSelfAttention::new(64, 4, 0.5, false, vs)?; // 50% dropout

    // Create input
    let input_data = vec![0.1f32; 640]; // 2 * 10 * 64
    let input = Var::from_slice(&input_data, (2, 10, 64), &device)?;

    // Forward pass
    let output = attention.forward(&input, false)?;

    // Compute loss
    let loss = output.sum_all()?;

    // Backward pass
    let grads = loss.backward()?;

    // Verify input gradient (dropout should scale gradients, not block them)
    let input_grad = grads
        .get(&input)
        .ok_or_else(|| MLError::ModelError("No gradient for input".to_string()))?;

    let grad_norm = compute_gradient_norm(input_grad)?;
    println!("Test 8 - Dropout gradient norm: {:.6}", grad_norm);

    assert!(
        grad_norm > 0.001,
        "Dropout should not completely block gradients: {:.6}",
        grad_norm
    );

    Ok(())
}

// ============================================================================
// TEST 9: TEMPERATURE SCALING GRADIENT FLOW
// ============================================================================

#[test]
fn test_temperature_scaling_gradient_flow() -> Result<(), MLError> {
    let device = Device::Cpu;

    // Create attention head with temperature
    let head = AttentionHead::new(64, 16, &device)?;

    // Create input
    let input_data = vec![0.1f32; 320]; // 2 * 10 * 64
    let input = Var::from_slice(&input_data, (2, 10, 64), &device)?;

    // Forward pass with temperature = 2.0
    let (output, _) = head.forward(&input, None, 2.0)?;

    // Compute loss
    let loss = output.sum_all()?;

    // Backward pass
    let grads = loss.backward()?;

    // Verify input gradient
    let input_grad = grads
        .get(&input)
        .ok_or_else(|| MLError::ModelError("No gradient for input".to_string()))?;

    let grad_norm = compute_gradient_norm(input_grad)?;
    println!("Test 9 - Temperature scaling gradient norm: {:.6}", grad_norm);

    assert!(
        grad_norm > 0.001,
        "Temperature scaling should preserve gradients: {:.6}",
        grad_norm
    );

    Ok(())
}

// ============================================================================
// TEST 10: GRADIENT CONSISTENCY ACROSS BATCH SIZES
// ============================================================================

#[test]
fn test_gradient_consistency_across_batch_sizes() -> Result<(), MLError> {
    let device = Device::Cpu;

    // Test with batch size 1
    let varmap1 = VarMap::new();
    let vs1 = VarBuilder::from_varmap(&varmap1, DType::F32, &device);
    let attention1 = TemporalSelfAttention::new(32, 2, 0.0, false, vs1)?;

    let input1_data = vec![0.1f32; 160]; // 1 * 10 * 16
    let input1 = Var::from_slice(&input1_data, (1, 10, 32), &device)?;
    let output1 = attention1.forward(&input1, false)?;
    let loss1 = output1.sum_all()?;
    let grads1 = loss1.backward()?;
    let grad1 = grads1.get(&input1).unwrap();
    let grad_norm1 = compute_gradient_norm(grad1)?;

    // Test with batch size 4
    let varmap2 = VarMap::new();
    let vs2 = VarBuilder::from_varmap(&varmap2, DType::F32, &device);
    let attention2 = TemporalSelfAttention::new(32, 2, 0.0, false, vs2)?;

    let input2_data = vec![0.1f32; 1280]; // 4 * 10 * 32
    let input2 = Var::from_slice(&input2_data, (4, 10, 32), &device)?;
    let output2 = attention2.forward(&input2, false)?;
    let loss2 = output2.sum_all()?;
    let grads2 = loss2.backward()?;
    let grad2 = grads2.get(&input2).unwrap();
    let grad_norm2 = compute_gradient_norm(grad2)?;

    println!(
        "Test 10 - Gradient norm (batch=1): {:.6}, (batch=4): {:.6}",
        grad_norm1, grad_norm2
    );

    // Both should have non-zero gradients
    assert!(grad_norm1 > 0.001, "Batch size 1 should have gradients");
    assert!(grad_norm2 > 0.001, "Batch size 4 should have gradients");

    Ok(())
}

// ============================================================================
// TEST 11: LONG SEQUENCE GRADIENT FLOW
// ============================================================================

#[test]
fn test_long_sequence_gradient_flow() -> Result<(), MLError> {
    let device = Device::Cpu;
    let varmap = VarMap::new();
    let vs = VarBuilder::from_varmap(&varmap, DType::F32, &device);

    // Create attention
    let attention = TemporalSelfAttention::new(64, 4, 0.1, false, vs)?;

    // Create long sequence input
    let seq_len = 100;
    let input_data = vec![0.1f32; 2 * seq_len * 64]; // batch=2, seq_len=100, hidden=64
    let input = Var::from_slice(&input_data, (2, seq_len, 64), &device)?;

    // Forward pass
    let output = attention.forward(&input, false)?;

    // Compute loss
    let loss = output.sum_all()?;

    // Backward pass
    let grads = loss.backward()?;

    // Verify input gradient
    let input_grad = grads
        .get(&input)
        .ok_or_else(|| MLError::ModelError("No gradient for input".to_string()))?;

    let grad_norm = compute_gradient_norm(input_grad)?;
    println!("Test 11 - Long sequence gradient norm: {:.6}", grad_norm);

    assert!(
        grad_norm > 0.001,
        "Long sequences should maintain gradient flow: {:.6}",
        grad_norm
    );

    Ok(())
}

// ============================================================================
// TEST 12: ALL ATTENTION HEADS RECEIVE GRADIENTS
// ============================================================================

#[test]
fn test_all_heads_receive_gradients() -> Result<(), MLError> {
    let device = Device::Cpu;
    let varmap = VarMap::new();
    let vs = VarBuilder::from_varmap(&varmap, DType::F32, &device);

    let num_heads = 8;
    let attention = TemporalSelfAttention::new(128, num_heads, 0.1, false, vs)?;

    // Create input
    let input_data = vec![0.1f32; 2560]; // 2 * 10 * 128
    let input = Var::from_slice(&input_data, (2, 10, 128), &device)?;

    // Forward pass
    let output = attention.forward(&input, false)?;

    // Compute loss
    let loss = output.sum_all()?;

    // Backward pass
    let grads = loss.backward()?;

    // Count how many variables have gradients
    let all_vars = varmap.all_vars();
    let mut vars_with_grads = 0;

    for var in &all_vars {
        if let Some(grad) = var.grad() {
            let grad_norm = compute_gradient_norm(&grad)?;
            if grad_norm > 1e-6 {
                vars_with_grads += 1;
            }
        }
    }

    println!(
        "Test 12 - Total variables: {}, with gradients: {}",
        all_vars.len(),
        vars_with_grads
    );

    // All heads should receive gradients (each head has query, key, value projections)
    // Plus output projection and layer norm parameters
    assert!(
        vars_with_grads > 0,
        "At least some attention parameters should have gradients"
    );

    Ok(())
}