foxhunt/ml/tests/tft_tests.rs

//! Comprehensive Tests for Temporal Fusion Transformer (TFT) Components
//!
//! Tests for:
//! 1. Temporal Attention - Multi-head self-attention with weight validation
//! 2. Variable Selection - Softmax gating with feature importance
//! 3. Gated Residual - GLU activation and skip connections
//! 4. Quantile Outputs - Multiple quantile predictions with ordering validation

#![allow(unused_crate_dependencies)]

use candle_core::{DType, Device, Tensor};
use candle_nn::VarBuilder;

use ml::tft::{
    GRNStack, GatedResidualNetwork, QuantileLayer, TemporalSelfAttention, VariableSelectionNetwork,
};
use ml::MLError;

// ============================================================================
// TEMPORAL ATTENTION TESTS
// ============================================================================

#[test]
fn test_attention_weights_sum_to_one() -> Result<(), MLError> {
    let device = Device::Cpu;
    let vs = VarBuilder::zeros(DType::F32, &device);

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

    // Create test input [batch_size=2, seq_len=5, hidden_dim=64]
    let input_data = vec![0.5f32; 640]; // 2 * 5 * 64
    let inputs = Tensor::from_slice(&input_data, (2, 5, 64), &device)?;

    let output = attention.forward(&inputs, true)?;

    // Output should maintain dimensions
    assert_eq!(output.dims(), &[2, 5, 64]);

    // Verify output is finite (no NaN or Inf)
    let output_data = output.flatten_all()?.to_vec1::<f32>()?;
    assert!(output_data.iter().all(|&x| x.is_finite()));

    Ok(())
}

#[test]
fn test_attention_causal_masking() -> Result<(), MLError> {
    let device = Device::Cpu;
    let vs = VarBuilder::zeros(DType::F32, &device);

    let attention = TemporalSelfAttention::new(32, 2, 0.0, false, vs)?;

    // Test that causal mask is properly applied
    let mask = attention.create_causal_mask(4)?;
    let mask_data = mask.to_vec2::<f32>()?;

    // Upper triangular should be -inf (masked)
    for i in 0..4 {
        for j in 0..4 {
            if j > i {
                assert!(
                    mask_data[i][j].is_infinite() && mask_data[i][j].is_sign_negative(),
                    "Position ({},{}) should be -inf, got {}",
                    i,
                    j,
                    mask_data[i][j]
                );
            } else {
                assert_eq!(
                    mask_data[i][j], 0.0,
                    "Position ({},{}) should be 0.0, got {}",
                    i, j, mask_data[i][j]
                );
            }
        }
    }

    Ok(())
}

#[test]
fn test_attention_positional_encoding() -> Result<(), MLError> {
    let device = Device::Cpu;
    let vs = VarBuilder::zeros(DType::F32, &device);

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

    // Test positional encoding at different sequence lengths
    let pos_enc_short = attention.positional_encoding.forward(10)?;
    let pos_enc_long = attention.positional_encoding.forward(50)?;

    assert_eq!(pos_enc_short.dims(), &[10, 64]);
    assert_eq!(pos_enc_long.dims(), &[50, 64]);

    // Verify sinusoidal pattern (different positions have different encodings)
    let short_data = pos_enc_short.to_vec2::<f32>()?;
    assert_ne!(
        short_data[0], short_data[1],
        "Different positions should have different encodings"
    );

    Ok(())
}

#[test]
fn test_attention_multi_head_output() -> Result<(), MLError> {
    let device = Device::Cpu;
    let vs = VarBuilder::zeros(DType::F32, &device);

    // Test with different head configurations
    for num_heads in [1, 2, 4, 8] {
        let hidden_dim = 64;
        assert_eq!(
            hidden_dim % num_heads,
            0,
            "Hidden dim must be divisible by num_heads"
        );

        let attention = TemporalSelfAttention::new(
            hidden_dim,
            num_heads,
            0.1,
            false,
            vs.pp(&format!("heads_{}", num_heads)),
        )?;

        let input_data = vec![0.5f32; 128]; // 2 * 64
        let inputs = Tensor::from_slice(&input_data, (2, 64), &device)?;
        let inputs_3d = inputs.unsqueeze(1)?; // [2, 1, 64]

        let output = attention.forward(&inputs_3d, false)?;
        assert_eq!(output.dims(), &[2, 1, 64]);
    }

    Ok(())
}

#[test]
fn test_attention_gradient_flow() -> Result<(), MLError> {
    let device = Device::Cpu;
    let vs = VarBuilder::zeros(DType::F32, &device);

    let attention = TemporalSelfAttention::new(32, 4, 0.1, false, vs)?;

    // Test with varying input magnitudes
    let small_input = Tensor::full(0.1f32, (2, 3, 32), &device)?;
    let large_input = Tensor::full(10.0f32, (2, 3, 32), &device)?;

    let small_output = attention.forward(&small_input, false)?;
    let large_output = attention.forward(&large_input, false)?;

    // Outputs should be different based on input magnitude
    let small_data = small_output.flatten_all()?.to_vec1::<f32>()?;
    let large_data = large_output.flatten_all()?.to_vec1::<f32>()?;

    let small_mean = small_data.iter().sum::<f32>() / small_data.len() as f32;
    let large_mean = large_data.iter().sum::<f32>() / large_data.len() as f32;

    assert_ne!(
        small_mean, large_mean,
        "Different inputs should produce different outputs"
    );

    Ok(())
}

// ============================================================================
// VARIABLE SELECTION TESTS
// ============================================================================

#[test]
fn test_variable_selection_gates_range() -> Result<(), MLError> {
    let device = Device::Cpu;
    let vs = VarBuilder::zeros(DType::F32, &device);

    let mut vsn = VariableSelectionNetwork::new(5, 32, vs.pp("test"))?;

    // Create test input
    let input_data = vec![1.0f32, 2.0, 3.0, 4.0, 5.0, 2.0, 3.0, 4.0, 5.0, 6.0];
    let inputs = Tensor::from_slice(&input_data, (2, 5), &device)?;

    let _output = vsn.forward(&inputs, None)?;

    // Check that importance scores are in valid range [0, 1] and sum to ~1
    let scores = vsn.get_importance_scores()?;
    assert_eq!(scores.len(), 5);

    for (i, &score) in scores.iter().enumerate() {
        assert!(
            score >= 0.0 && score <= 1.0,
            "Score {} at index {} is out of range [0,1]",
            score,
            i
        );
    }

    // Should sum to approximately 1.0 (softmax normalization)
    let sum: f64 = scores.iter().sum();
    assert!(
        (sum - 1.0).abs() < 1e-6,
        "Importance scores should sum to 1.0, got {}",
        sum
    );

    Ok(())
}

#[test]
fn test_variable_selection_feature_importance() -> Result<(), MLError> {
    let device = Device::Cpu;
    let vs = VarBuilder::zeros(DType::F32, &device);

    let mut vsn = VariableSelectionNetwork::new(10, 64, vs.pp("test"))?;

    // Create input with varying magnitudes to encourage selection
    let mut input_data = Vec::new();
    for _batch in 0..2 {
        for i in 0..10 {
            input_data.push((i as f32) * 0.5); // Different magnitudes
        }
    }
    let inputs = Tensor::from_slice(&input_data, (2, 10), &device)?;

    let _output = vsn.forward(&inputs, None)?;

    // Get top features
    let top_features = vsn.get_top_features(3);
    assert_eq!(top_features.len(), 3);

    // Verify features are sorted by importance (descending)
    for i in 1..top_features.len() {
        assert!(
            top_features[i - 1].1 >= top_features[i].1,
            "Features should be sorted by importance"
        );
    }

    // Verify importance scores are valid
    for (idx, score) in &top_features {
        assert!(*idx < 10, "Feature index {} out of range", idx);
        assert!(
            *score >= 0.0 && *score <= 1.0,
            "Score {} out of range",
            score
        );
    }

    Ok(())
}

#[test]
fn test_variable_selection_with_context() -> Result<(), MLError> {
    let device = Device::Cpu;
    let vs = VarBuilder::zeros(DType::F32, &device);

    let mut vsn = VariableSelectionNetwork::new(5, 32, vs.pp("test"))?;

    let input_data = vec![1.0f32; 10]; // 2 * 5
    let inputs = Tensor::from_slice(&input_data, (2, 5), &device)?;

    let context_data = vec![0.5f32; 64]; // 2 * 32
    let context = Tensor::from_slice(&context_data, (2, 32), &device)?;

    // Forward without context
    let output_no_ctx = vsn.forward(&inputs, None)?;

    // Reset for fair comparison (create new network with same config)
    let mut vsn2 = VariableSelectionNetwork::new(5, 32, vs.pp("test2"))?;
    let output_with_ctx = vsn2.forward(&inputs, Some(&context))?;

    // Both should produce valid outputs
    assert_eq!(output_no_ctx.dims(), &[2, 1, 32]);
    assert_eq!(output_with_ctx.dims(), &[2, 1, 32]);

    Ok(())
}

#[test]
fn test_variable_selection_3d_input() -> Result<(), MLError> {
    let device = Device::Cpu;
    let vs = VarBuilder::zeros(DType::F32, &device);

    let mut vsn = VariableSelectionNetwork::new(4, 24, vs.pp("test"))?;

    // Test 3D input [batch_size=2, seq_len=3, input_size=4]
    let input_data = vec![1.0f32; 24]; // 2 * 3 * 4
    let inputs = Tensor::from_slice(&input_data, (2, 3, 4), &device)?;

    let output = vsn.forward(&inputs, None)?;

    // Should produce [batch_size=2, seq_len=3, hidden_size=24]
    assert_eq!(output.dims(), &[2, 3, 24]);

    Ok(())
}

// ============================================================================
// GATED RESIDUAL NETWORK TESTS
// ============================================================================

#[test]
fn test_grn_skip_connection() -> Result<(), MLError> {
    let device = Device::Cpu;
    let vs = VarBuilder::zeros(DType::F32, &device);

    // Test skip connection when input/output dims are same
    let grn_same = GatedResidualNetwork::new(32, 32, vs.pp("same"))?;

    let input_data = vec![1.0f32; 64]; // 2 * 32
    let inputs = Tensor::from_slice(&input_data, (2, 32), &device)?;

    let output = grn_same.forward(&inputs, None)?;
    assert_eq!(output.dims(), &[2, 32]);

    // Test skip connection when input/output dims differ
    let grn_diff = GatedResidualNetwork::new(64, 32, vs.pp("diff"))?;
    let input_data_diff = vec![1.0f32; 128]; // 2 * 64
    let inputs_diff = Tensor::from_slice(&input_data_diff, (2, 64), &device)?;

    let output_diff = grn_diff.forward(&inputs_diff, None)?;
    assert_eq!(output_diff.dims(), &[2, 32]);

    Ok(())
}

#[test]
fn test_grn_glu_activation() -> Result<(), MLError> {
    let device = Device::Cpu;
    let vs = VarBuilder::zeros(DType::F32, &device);

    let grn = GatedResidualNetwork::new(32, 32, vs.pp("test"))?;

    // Test with different input magnitudes to verify gating
    let zero_input = Tensor::zeros((2, 32), DType::F32, &device)?;
    let nonzero_input = Tensor::ones((2, 32), DType::F32, &device)?;

    let zero_output = grn.forward(&zero_input, None)?;
    let nonzero_output = grn.forward(&nonzero_input, None)?;

    // Outputs should differ based on input
    let zero_data = zero_output.flatten_all()?.to_vec1::<f32>()?;
    let nonzero_data = nonzero_output.flatten_all()?.to_vec1::<f32>()?;

    let zero_mean = zero_data.iter().sum::<f32>() / zero_data.len() as f32;
    let nonzero_mean = nonzero_data.iter().sum::<f32>() / nonzero_data.len() as f32;

    // GLU gating should produce different outputs for different inputs
    assert_ne!(
        zero_mean, nonzero_mean,
        "GLU should produce different outputs for different inputs"
    );

    Ok(())
}

#[test]
fn test_grn_context_integration() -> Result<(), MLError> {
    let device = Device::Cpu;
    let vs = VarBuilder::zeros(DType::F32, &device);

    let grn = GatedResidualNetwork::new(32, 32, vs.pp("test"))?;

    let input_data = vec![1.0f32; 64]; // 2 * 32
    let inputs = Tensor::from_slice(&input_data, (2, 32), &device)?;

    let context_data = vec![2.0f32; 64]; // 2 * 32
    let context = Tensor::from_slice(&context_data, (2, 32), &device)?;

    // Test without context
    let output_no_ctx = grn.forward(&inputs, None)?;

    // Test with context
    let output_with_ctx = grn.forward(&inputs, Some(&context))?;

    // Both should produce valid outputs
    assert_eq!(output_no_ctx.dims(), &[2, 32]);
    assert_eq!(output_with_ctx.dims(), &[2, 32]);

    // Outputs should differ when context is provided
    let no_ctx_data = output_no_ctx.flatten_all()?.to_vec1::<f32>()?;
    let with_ctx_data = output_with_ctx.flatten_all()?.to_vec1::<f32>()?;

    // At least some values should differ
    let differences = no_ctx_data
        .iter()
        .zip(with_ctx_data.iter())
        .filter(|(a, b)| (*a - *b).abs() > 1e-6)
        .count();

    assert!(
        differences > 0,
        "Context should affect output (found {} different values)",
        differences
    );

    Ok(())
}

#[test]
fn test_grn_stack_depth() -> Result<(), MLError> {
    let device = Device::Cpu;
    let vs = VarBuilder::zeros(DType::F32, &device);

    // Test different stack depths
    for num_layers in [1, 2, 3, 5] {
        let stack = GRNStack::new(
            64,
            32,
            16,
            num_layers,
            vs.pp(&format!("stack_{}", num_layers)),
        )?;

        assert_eq!(stack.num_layers, num_layers);

        let input_data = vec![1.0f32; 128]; // 2 * 64
        let inputs = Tensor::from_slice(&input_data, (2, 64), &device)?;

        let output = stack.forward(&inputs, None)?;

        // Final output should match final layer output dim
        assert_eq!(output.dims(), &[2, 16]);
    }

    Ok(())
}

#[test]
fn test_grn_gradient_flow() -> Result<(), MLError> {
    let device = Device::Cpu;
    let vs = VarBuilder::zeros(DType::F32, &device);

    let grn = GatedResidualNetwork::new(32, 32, vs.pp("test"))?;

    // Test with varying input scales
    let scales = [0.1f32, 1.0, 10.0];
    let mut outputs = Vec::new();

    for &scale in &scales {
        let input = Tensor::full(scale, (2, 32), &device)?;
        let output = grn.forward(&input, None)?;
        outputs.push(output);
    }

    // Verify that different scales produce different outputs (gradient flow)
    for i in 0..outputs.len() - 1 {
        let out1 = outputs[i].flatten_all()?.to_vec1::<f32>()?;
        let out2 = outputs[i + 1].flatten_all()?.to_vec1::<f32>()?;

        let mean1 = out1.iter().sum::<f32>() / out1.len() as f32;
        let mean2 = out2.iter().sum::<f32>() / out2.len() as f32;

        assert_ne!(
            mean1, mean2,
            "Different input scales should produce different outputs"
        );
    }

    Ok(())
}

// ============================================================================
// QUANTILE OUTPUT TESTS
// ============================================================================

#[test]
fn test_quantile_ordering_validation() -> Result<(), MLError> {
    let device = Device::Cpu;
    let vs = VarBuilder::zeros(DType::F32, &device);

    let quantile_layer = QuantileLayer::new(32, 5, 9, vs.pp("test"))?;

    // Create test input
    let input_data = vec![1.0f32; 64]; // 2 * 32
    let inputs = Tensor::from_slice(&input_data, (2, 32), &device)?;

    let output = quantile_layer.forward(&inputs)?;

    // Output shape should be [batch_size=2, prediction_horizon=5, num_quantiles=9]
    assert_eq!(output.dims(), &[2, 5, 9]);

    // Verify quantile ordering: q_i <= q_{i+1} for all i
    let output_data = output.to_vec3::<f32>()?;

    for batch in 0..2 {
        for horizon in 0..5 {
            let quantiles = &output_data[batch][horizon];

            // Check monotonic ordering
            for i in 1..quantiles.len() {
                assert!(
                    quantiles[i] >= quantiles[i - 1],
                    "Quantiles not monotonic at batch={}, horizon={}, q[{}]={} < q[{}]={}",
                    batch,
                    horizon,
                    i,
                    quantiles[i],
                    i - 1,
                    quantiles[i - 1]
                );
            }
        }
    }

    Ok(())
}

#[test]
fn test_quantile_levels_correct() -> Result<(), MLError> {
    let device = Device::Cpu;
    let vs = VarBuilder::zeros(DType::F32, &device);

    let quantile_layer = QuantileLayer::new(32, 5, 9, vs.pp("test"))?;
    let levels = quantile_layer.get_quantile_levels();

    // Should generate 9 evenly spaced quantile levels
    assert_eq!(levels.len(), 9);

    // Check that levels are approximately [0.1, 0.2, ..., 0.9]
    for (i, &level) in levels.iter().enumerate() {
        let expected = (i + 1) as f64 / 10.0; // 0.1, 0.2, ..., 0.9
        assert!(
            (level - expected).abs() < 0.01,
            "Quantile level {} should be approximately {}, got {}",
            i,
            expected,
            level
        );
    }

    // Verify monotonic increase
    for i in 1..levels.len() {
        assert!(
            levels[i] > levels[i - 1],
            "Quantile levels should be monotonically increasing"
        );
    }

    Ok(())
}

#[test]
fn test_quantile_prediction_intervals() -> Result<(), MLError> {
    let device = Device::Cpu;
    let vs = VarBuilder::zeros(DType::F32, &device);

    let quantile_layer = QuantileLayer::new(16, 3, 9, vs.pp("test"))?;

    // Create mock quantile predictions with known ordering
    let mut quantile_data = Vec::new();
    for _batch in 0..2 {
        for _horizon in 0..3 {
            for q in 1..=9 {
                quantile_data.push(q as f32); // 1.0, 2.0, ..., 9.0
            }
        }
    }
    let quantiles = Tensor::from_slice(&quantile_data, (2, 3, 9), &device)?;

    // Test different confidence levels
    for &confidence in &[0.50, 0.80, 0.90, 0.95] {
        let (lower, upper) = quantile_layer.get_prediction_intervals(&quantiles, confidence)?;

        assert_eq!(lower.dims(), &[2, 3]);
        assert_eq!(upper.dims(), &[2, 3]);

        // Upper bound should always be >= lower bound
        let lower_data = lower.to_vec2::<f32>()?;
        let upper_data = upper.to_vec2::<f32>()?;

        for batch in 0..2 {
            for horizon in 0..3 {
                assert!(
                    upper_data[batch][horizon] >= lower_data[batch][horizon],
                    "Upper bound {} should be >= lower bound {} for confidence {}",
                    upper_data[batch][horizon],
                    lower_data[batch][horizon],
                    confidence
                );
            }
        }
    }

    Ok(())
}

#[test]
fn test_quantile_loss_computation() -> Result<(), MLError> {
    let device = Device::Cpu;
    let vs = VarBuilder::zeros(DType::F32, &device);

    let quantile_layer = QuantileLayer::new(16, 2, 3, vs.pp("test"))?;

    // Create predictions [batch=2, horizon=2, quantiles=3]
    let pred_data = vec![
        1.0f32, 2.0, 3.0, // batch 0, horizon 0
        1.5, 2.5, 3.5, // batch 0, horizon 1
        2.0, 3.0, 4.0, // batch 1, horizon 0
        2.5, 3.5, 4.5, // batch 1, horizon 1
    ];
    let predictions = Tensor::from_slice(&pred_data, (2, 2, 3), &device)?;

    // Create targets [batch=2, horizon=2]
    let target_data = vec![2.0f32, 2.5, 3.0, 3.5];
    let targets = Tensor::from_slice(&target_data, (2, 2), &device)?;

    let loss = quantile_layer.quantile_loss(&predictions, &targets)?;

    // Loss should be a scalar
    assert_eq!(loss.dims(), &[] as &[usize]);

    // Loss should be non-negative
    let loss_value = loss.to_vec0::<f32>()?;
    assert!(loss_value >= 0.0, "Quantile loss should be non-negative");

    // Loss should be finite
    assert!(loss_value.is_finite(), "Quantile loss should be finite");

    Ok(())
}

#[test]
fn test_quantile_loss_symmetry() -> Result<(), MLError> {
    let device = Device::Cpu;
    let vs = VarBuilder::zeros(DType::F32, &device);

    let quantile_layer = QuantileLayer::new(16, 2, 5, vs.pp("test"))?;

    // Create symmetric predictions around target
    let pred_data = vec![
        1.0f32, 1.5, 2.0, 2.5, 3.0, // quantiles for horizon 0
        1.0, 1.5, 2.0, 2.5, 3.0, // quantiles for horizon 1
    ];
    let predictions = Tensor::from_slice(&pred_data, (1, 2, 5), &device)?;

    // Target at median (2.0)
    let target_data = vec![2.0f32, 2.0];
    let targets = Tensor::from_slice(&target_data, (1, 2), &device)?;

    let loss = quantile_layer.quantile_loss(&predictions, &targets)?;
    let loss_value = loss.to_vec0::<f32>()?;

    // Loss should be relatively small when target is at median
    assert!(
        loss_value < 1.0,
        "Loss should be small when predictions are symmetric around target"
    );

    Ok(())
}

#[test]
fn test_quantile_3d_input_handling() -> Result<(), MLError> {
    let device = Device::Cpu;
    let vs = VarBuilder::zeros(DType::F32, &device);

    let quantile_layer = QuantileLayer::new(16, 3, 5, vs.pp("test"))?;

    // Test 3D input [batch_size=2, seq_len=10, hidden_dim=16]
    let input_data = vec![1.0f32; 320]; // 2 * 10 * 16
    let inputs = Tensor::from_slice(&input_data, (2, 10, 16), &device)?;

    let output = quantile_layer.forward(&inputs)?;

    // Should use last time step and produce [batch_size=2, horizon=3, quantiles=5]
    assert_eq!(output.dims(), &[2, 3, 5]);

    // Verify quantile ordering for 3D input
    let output_data = output.to_vec3::<f32>()?;

    for batch in 0..2 {
        for horizon in 0..3 {
            let quantiles = &output_data[batch][horizon];
            for i in 1..quantiles.len() {
                assert!(
                    quantiles[i] >= quantiles[i - 1],
                    "Quantiles should be monotonic even with 3D input"
                );
            }
        }
    }

    Ok(())
}

// ============================================================================
// INTEGRATION TESTS
// ============================================================================

#[test]
fn test_tft_component_integration() -> Result<(), MLError> {
    let device = Device::Cpu;
    let vs = VarBuilder::zeros(DType::F32, &device);

    // Create all components
    let mut vsn = VariableSelectionNetwork::new(10, 32, vs.pp("vsn"))?;
    let grn = GatedResidualNetwork::new(32, 32, vs.pp("grn"))?;
    let attention = TemporalSelfAttention::new(32, 4, 0.1, false, vs.pp("attn"))?;
    let quantile = QuantileLayer::new(32, 5, 7, vs.pp("quant"))?;

    // Simulate TFT pipeline
    // 1. Variable selection
    let input_data = vec![1.0f32; 20]; // 2 * 10
    let inputs = Tensor::from_slice(&input_data, (2, 10), &device)?;
    let selected = vsn.forward(&inputs, None)?;

    // 2. Gated residual
    let selected_2d = selected.squeeze(1)?; // [2, 32]
    let encoded = grn.forward(&selected_2d, None)?;

    // 3. Attention
    let encoded_3d = encoded.unsqueeze(1)?; // [2, 1, 32]
    let attended = attention.forward(&encoded_3d, false)?;

    // 4. Quantile output
    let attended_2d = attended.squeeze(1)?; // [2, 32]
    let quantiles = quantile.forward(&attended_2d)?;

    // Verify final output shape
    assert_eq!(quantiles.dims(), &[2, 5, 7]);

    // Verify quantile ordering in integrated pipeline
    let quantile_data = quantiles.to_vec3::<f32>()?;
    for batch in 0..2 {
        for horizon in 0..5 {
            let q = &quantile_data[batch][horizon];
            for i in 1..q.len() {
                assert!(
                    q[i] >= q[i - 1],
                    "Quantile ordering preserved through pipeline"
                );
            }
        }
    }

    Ok(())
}

#[test]
fn test_attention_weight_normalization() -> Result<(), MLError> {
    let device = Device::Cpu;
    let vs = VarBuilder::zeros(DType::F32, &device);

    let attention = TemporalSelfAttention::new(64, 8, 0.0, false, vs)?;

    // Test multiple batch sizes and sequence lengths
    for (batch_size, seq_len) in [(1, 5), (2, 10), (4, 20)] {
        let input_size = batch_size * seq_len * 64;
        let input_data = vec![0.5f32; input_size];
        let inputs = Tensor::from_slice(&input_data, (batch_size, seq_len, 64), &device)?;

        let output = attention.forward(&inputs, true)?;

        // Verify output shape and values are valid
        assert_eq!(output.dims(), &[batch_size, seq_len, 64]);

        let output_vec = output.flatten_all()?.to_vec1::<f32>()?;
        assert!(
            output_vec.iter().all(|&x| x.is_finite()),
            "All attention outputs should be finite"
        );
    }

    Ok(())
}

#[test]
fn test_variable_selection_consistency() -> Result<(), MLError> {
    let device = Device::Cpu;
    let vs = VarBuilder::zeros(DType::F32, &device);

    let mut vsn = VariableSelectionNetwork::new(8, 48, vs.pp("test"))?;

    // Same input should produce consistent importance scores
    let input_data = vec![1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0];
    let inputs = Tensor::from_slice(&input_data, (1, 8), &device)?;

    let _output1 = vsn.forward(&inputs, None)?;
    let scores1 = vsn.get_importance_scores()?;

    let _output2 = vsn.forward(&inputs, None)?;
    let scores2 = vsn.get_importance_scores()?;

    // Scores should be identical for same input
    for (i, (s1, s2)) in scores1.into_iter().zip(scores2.into_iter()).enumerate() {
        assert!(
            (s1 - s2).abs() < 1e-6,
            "Score {} differs: {} vs {}",
            i,
            s1,
            s2
        );
    }

    Ok(())
}