foxhunt/ml/tests/test_tft_cuda_layernorm.rs

//! Integration test for TFT with CUDA-compatible layer normalization
//!
//! This test validates that TFT model can perform forward passes
//! with the new manual CUDA layer normalization implementation.

use anyhow::Result;
use candle_core::{DType, Device, Tensor};
use ml::tft::{TFTConfig, TemporalFusionTransformer};

#[test]
fn test_tft_forward_pass_with_cuda_layernorm() -> Result<()> {
    // Create small TFT config for testing
    let config = TFTConfig {
        input_dim: 10,
        hidden_dim: 32,
        num_heads: 4,
        num_layers: 2,
        prediction_horizon: 5,
        sequence_length: 20,
        num_quantiles: 5,
        num_static_features: 2,
        num_known_features: 3,
        num_unknown_features: 5,
        ..Default::default()
    };

    // Create TFT model (automatically uses CUDA if available)
    let mut tft = TemporalFusionTransformer::new(config.clone())?;

    // Get device (CUDA if available, CPU otherwise)
    let device = Device::cuda_if_available(0).unwrap_or(Device::Cpu);
    println!("Testing on device: {:?}", device);

    // Create test inputs
    let batch_size = 2;

    // Static features [batch_size, num_static_features]
    let static_features =
        Tensor::randn(0f32, 1.0, (batch_size, config.num_static_features), &device)?;

    // Historical features [batch_size, sequence_length, num_unknown_features]
    let historical_features = Tensor::randn(
        0f32,
        1.0,
        (
            batch_size,
            config.sequence_length,
            config.num_unknown_features,
        ),
        &device,
    )?;

    // Future features [batch_size, prediction_horizon, num_known_features]
    let future_features = Tensor::randn(
        0f32,
        1.0,
        (
            batch_size,
            config.prediction_horizon,
            config.num_known_features,
        ),
        &device,
    )?;

    // Perform forward pass
    let start = std::time::Instant::now();
    let output = tft.forward(&static_features, &historical_features, &future_features)?;
    let duration = start.elapsed();

    println!("Forward pass completed in {:?}", duration);

    // Validate output shape
    // Expected: [batch_size, prediction_horizon, num_quantiles]
    let expected_shape = &[batch_size, config.prediction_horizon, config.num_quantiles];
    assert_eq!(
        output.dims(),
        expected_shape,
        "Output shape mismatch. Expected {:?}, got {:?}",
        expected_shape,
        output.dims()
    );

    // Validate output values (no NaN, no Inf)
    let output_vec = output.flatten_all()?.to_vec1::<f32>()?;
    let has_nan = output_vec.iter().any(|&x| x.is_nan());
    let has_inf = output_vec.iter().any(|&x| x.is_infinite());

    assert!(!has_nan, "Output contains NaN values");
    assert!(!has_inf, "Output contains Inf values");

    println!("✅ TFT forward pass successful with CUDA layer normalization");
    println!("   Output shape: {:?}", output.dims());
    println!(
        "   Output range: [{:.4}, {:.4}]",
        output_vec.iter().cloned().fold(f32::INFINITY, f32::min),
        output_vec.iter().cloned().fold(f32::NEG_INFINITY, f32::max)
    );

    Ok(())
}

#[test]
fn test_tft_grn_with_cuda_layernorm() -> Result<()> {
    use candle_nn::VarBuilder;
    use ml::tft::gated_residual::GatedResidualNetwork;

    let device = Device::cuda_if_available(0).unwrap_or(Device::Cpu);
    println!("Testing GRN on device: {:?}", device);

    let vs = VarBuilder::zeros(DType::F32, &device);
    let grn = GatedResidualNetwork::new(64, 32, vs.pp("test"))?;

    // Create test input [batch_size=2, hidden_dim=64]
    let input = Tensor::randn(0f32, 1.0, (2, 64), &device)?;

    // Forward pass (uses CudaLayerNorm internally)
    let output = grn.forward(&input, None)?;

    // Validate output
    assert_eq!(output.dims(), &[2, 32]);

    let output_vec = output.flatten_all()?.to_vec1::<f32>()?;
    let has_nan = output_vec.iter().any(|&x| x.is_nan());
    let has_inf = output_vec.iter().any(|&x| x.is_infinite());

    assert!(!has_nan, "GRN output contains NaN values");
    assert!(!has_inf, "GRN output contains Inf values");

    println!("✅ GRN forward pass successful with CUDA layer normalization");
    println!("   Output shape: {:?}", output.dims());

    Ok(())
}

#[test]
fn test_tft_attention_with_cuda_layernorm() -> Result<()> {
    use candle_nn::VarBuilder;
    use ml::tft::temporal_attention::TemporalSelfAttention;

    let device = Device::cuda_if_available(0).unwrap_or(Device::Cpu);
    println!("Testing Temporal Attention on device: {:?}", device);

    let vs = VarBuilder::zeros(DType::F32, &device);
    let attention = TemporalSelfAttention::new(
        256,  // hidden_dim
        8,    // num_heads
        0.1,  // dropout_rate
        true, // use_flash_attention
        vs,
    )?;

    // Create test input [batch_size=2, seq_len=10, hidden_dim=256]
    let input = Tensor::randn(0f32, 1.0, (2, 10, 256), &device)?;

    // Forward pass (uses CudaLayerNorm internally)
    let output = attention.forward(&input, true)?;

    // Validate output
    assert_eq!(output.dims(), &[2, 10, 256]);

    let output_vec = output.flatten_all()?.to_vec1::<f32>()?;
    let has_nan = output_vec.iter().any(|&x| x.is_nan());
    let has_inf = output_vec.iter().any(|&x| x.is_infinite());

    assert!(!has_nan, "Attention output contains NaN values");
    assert!(!has_inf, "Attention output contains Inf values");

    println!("✅ Temporal Attention forward pass successful with CUDA layer normalization");
    println!("   Output shape: {:?}", output.dims());

    Ok(())
}

#[test]
fn test_tft_batch_processing() -> Result<()> {
    // Test with various batch sizes to ensure layer norm handles broadcasting correctly
    let config = TFTConfig {
        input_dim: 10,
        hidden_dim: 32,
        num_heads: 4,
        num_layers: 1,
        prediction_horizon: 3,
        sequence_length: 10,
        num_quantiles: 3,
        num_static_features: 2,
        num_known_features: 2,
        num_unknown_features: 6  // 2 + 2 + 6 = 10 (fixed feature count mismatch),
        ..Default::default(),
    };

    let device = Device::cuda_if_available(0).unwrap_or(Device::Cpu);
    let mut tft = TemporalFusionTransformer::new(config.clone())?;

    for batch_size in [1, 2, 4, 8] {
        let static_features =
            Tensor::randn(0f32, 1.0, (batch_size, config.num_static_features), &device)?;

        let historical_features = Tensor::randn(
            0f32,
            1.0,
            (
                batch_size,
                config.sequence_length,
                config.num_unknown_features,
            ),
            &device,
        )?;

        let future_features = Tensor::randn(
            0f32,
            1.0,
            (
                batch_size,
                config.prediction_horizon,
                config.num_known_features,
            ),
            &device,
        )?;

        let output = tft.forward(&static_features, &historical_features, &future_features)?;

        assert_eq!(
            output.dims(),
            &[batch_size, config.prediction_horizon, config.num_quantiles],
            "Batch size {} failed",
            batch_size
        );

        println!("✅ Batch size {} processed successfully", batch_size);
    }

    Ok(())
}