foxhunt/crates/ml/tests/test_quantile_output_standalone.rs

/// Standalone test for forward_quantile_output method
///
/// Tests the core quantile output layer in isolation
use candle_core::{Device, Tensor};
use ml::memory_optimization::quantization::{QuantizationConfig, QuantizationType, Quantizer};
use ml::tft::{QuantizedTemporalFusionTransformer, TFTConfig};
use ml::MLError;

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

    // Create TFT config
    let mut config = TFTConfig::default();
    config.num_quantiles = 3;
    config.prediction_horizon = 10;
    config.hidden_dim = 256;

    let tft = QuantizedTemporalFusionTransformer::new_with_device(config.clone(), device.clone())?;

    // Create test inputs
    let batch_size = 2;

    // Decoder output: [batch, horizon, hidden_dim]
    let decoder_output = Tensor::randn(
        0f32,
        1.0,
        (batch_size, config.prediction_horizon, config.hidden_dim),
        &device,
    )?;

    // Output projection weights: [hidden_dim, num_quantiles]
    let weight_data = Tensor::randn(
        0f32,
        0.01f32,
        (config.hidden_dim, config.num_quantiles),
        &device,
    )?;

    // Quantize the weights
    let mut quantizer = Quantizer::new(
        QuantizationConfig {
            quant_type: QuantizationType::Int8,
            per_channel: false,
            symmetric: true,
            calibration_samples: None,
        },
        device.clone(),
    );

    let quantized_weights = quantizer.quantize_tensor(&weight_data, "output_projection")?;

    // Test forward_quantile_output
    let output = tft.forward_quantile_output(&decoder_output, &quantized_weights)?;

    // Validate output shape: [batch=2, horizon=10, quantiles=3]
    assert_eq!(
        output.dims(),
        &[batch_size, config.prediction_horizon, config.num_quantiles],
        "Output shape mismatch"
    );

    // Validate no NaN/Inf
    let output_data = output.flatten_all()?.to_vec1::<f32>()?;
    assert!(
        output_data.iter().all(|&x| x.is_finite()),
        "Output contains NaN or Inf"
    );

    // Test that output values are within reasonable range
    let max_val = output_data.iter().fold(f32::NEG_INFINITY, |a, &b| a.max(b));
    let min_val = output_data.iter().fold(f32::INFINITY, |a, &b| a.min(b));
    assert!(
        max_val.abs() < 100.0 && min_val.abs() < 100.0,
        "Output values out of reasonable range: min={}, max={}",
        min_val,
        max_val
    );

    println!("✅ forward_quantile_output test passed!");
    println!("   Output shape: {:?}", output.dims());
    println!("   Output range: [{:.4}, {:.4}]", min_val, max_val);

    Ok(())
}

#[test]
fn test_forward_quantile_output_invalid_dims() {
    let device = Device::Cpu;
    let config = TFTConfig::default();
    let tft = QuantizedTemporalFusionTransformer::new_with_device(config.clone(), device.clone())
        .expect("Failed to create TFT");

    // Create invalid 2D input (should be 3D)
    let invalid_input =
        Tensor::zeros((2, 256), candle_core::DType::F32, &device).expect("Failed to create tensor");

    let weight_data = Tensor::zeros((256, 3), candle_core::DType::F32, &device)
        .expect("Failed to create weights");

    let mut quantizer = Quantizer::new(
        QuantizationConfig {
            quant_type: QuantizationType::Int8,
            per_channel: false,
            symmetric: true,
            calibration_samples: None,
        },
        device.clone(),
    );

    let quantized_weights = quantizer
        .quantize_tensor(&weight_data, "test_weights")
        .expect("Failed to quantize");

    let result = tft.forward_quantile_output(&invalid_input, &quantized_weights);
    assert!(result.is_err(), "Should reject 2D input");

    match result {
        Err(MLError::InvalidInput(msg)) => {
            assert!(
                msg.contains("3 dimensions"),
                "Error message should mention 3 dimensions: {}",
                msg
            );
        },
        _ => panic!("Expected InvalidInput error"),
    }
}