foxhunt/ml/examples/test_tft_fp32_varmap.rs

//! Test FP32 TFT model parameter count and forward pass functionality
//!
//! This example verifies that the non-quantized TFT model:
//! 1. Has real trainable parameters in its VarMap
//! 2. Can execute forward passes successfully
//! 3. Produces non-dummy output tensors
//! 4. Integrates correctly with the AdamW optimizer

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

fn main() -> Result<(), Box<dyn std::error::Error>> {
    println!("=== FP32 TFT Model Parameter Analysis ===\n");

    let config = TFTConfig {
        input_dim: 225,
        hidden_dim: 256,
        num_heads: 8,
        num_layers: 2,
        prediction_horizon: 10,
        sequence_length: 60,
        num_quantiles: 3,
        num_static_features: 5,
        num_known_features: 10,
        num_unknown_features: 210,
        ..Default::default()
    };

    println!("Creating FP32 TFT model with config:");
    println!("  Input dim: {}", config.input_dim);
    println!("  Hidden dim: {}", config.hidden_dim);
    println!("  Num layers: {}", config.num_layers);
    println!("  Num heads: {}", config.num_heads);

    let device = Device::Cpu;
    let model = TemporalFusionTransformer::new_with_device(config, device)?;

    println!("\n=== VarMap Analysis ===");
    let varmap = model.get_varmap();
    let all_vars = varmap.all_vars();

    println!("Total number of parameter tensors: {}", all_vars.len());

    let mut total_params = 0usize;
    for (i, var) in all_vars.iter().enumerate() {
        let shape = var.shape();
        let param_count: usize = shape.dims().iter().product();
        total_params += param_count;

        if i < 10 {
            // Show first 10 parameters
            println!(
                "  Param {}: shape {:?}, count {}",
                i,
                shape.dims(),
                param_count
            );
        }
    }

    if all_vars.len() > 10 {
        println!("  ... ({} more parameters)", all_vars.len() - 10);
    }

    println!("\nTotal trainable parameters: {}", total_params);
    println!(
        "Estimated FP32 memory (4 bytes/param): {:.2} MB",
        (total_params * 4) as f64 / 1_048_576.0
    );

    println!("\n=== Forward Pass Test ===");

    let batch_size = 2;
    let seq_len = 60;
    let horizon = 10;

    let static_features = Tensor::zeros(&[batch_size, 5], candle_core::DType::F32, &Device::Cpu)?;
    let historical_features = Tensor::zeros(
        &[batch_size, seq_len, 210],
        candle_core::DType::F32,
        &Device::Cpu,
    )?;
    let future_features = Tensor::zeros(
        &[batch_size, horizon, 10],
        candle_core::DType::F32,
        &Device::Cpu,
    )?;

    let mut model_mut = model;
    let output = model_mut.forward(&static_features, &historical_features, &future_features)?;

    println!("Input shapes:");
    println!("  Static: [batch={}, features=5]", batch_size);
    println!(
        "  Historical: [batch={}, seq={}, features=210]",
        batch_size, seq_len
    );
    println!(
        "  Future: [batch={}, horizon={}, features=10]",
        batch_size, horizon
    );

    println!("\nOutput shape: {:?}", output.shape());
    println!(
        "Expected: [batch={}, horizon={}, quantiles=3]",
        batch_size, horizon
    );

    // Check if output contains actual computed values (not zeros/NaN)
    let output_data = output.flatten_all()?.to_vec1::<f32>()?;
    let has_nonzero = output_data.iter().any(|&x| x.abs() > 1e-10);
    let has_nan = output_data.iter().any(|&x| x.is_nan());
    let has_inf = output_data.iter().any(|&x| x.is_infinite());

    println!("\nOutput validation:");
    println!("  Has non-zero values: {}", has_nonzero);
    println!("  Has NaN values: {}", has_nan);
    println!("  Has Inf values: {}", has_inf);

    // Sample a few output values
    println!("\nSample output values (first 5):");
    for (i, &val) in output_data.iter().take(5).enumerate() {
        println!("  output[{}] = {:.6e}", i, val);
    }

    println!("\n=== Optimizer Integration Test ===");
    use candle_nn::Optimizer;
    use candle_optimisers::adam::{Adam, ParamsAdam};

    let optimizer_params = ParamsAdam {
        lr: 1e-3,
        beta_1: 0.9,
        beta_2: 0.999,
        eps: 1e-8,
        weight_decay: None,
        amsgrad: false,
    };

    let mut optimizer = Adam::new(all_vars.clone(), optimizer_params)?;

    println!(
        "AdamW optimizer created with {} parameter groups",
        all_vars.len()
    );

    // Simulate a backward pass
    let target = Tensor::zeros(
        &[batch_size, horizon, 3],
        candle_core::DType::F32,
        &Device::Cpu,
    )?;
    let loss = ((output - target)?.sqr()?.sum_all())?;
    let loss_value = loss.to_vec0::<f32>()?;

    println!("Computed loss: {:.6}", loss_value);

    // Try optimizer step
    let step_result = optimizer.backward_step(&loss);
    println!(
        "Optimizer step: {}",
        if step_result.is_ok() {
            "✅ Success"
        } else {
            "❌ Failed"
        }
    );

    println!("\n=== Verdict ===");
    if all_vars.is_empty() {
        println!("❌ BROKEN: VarMap is EMPTY - no trainable parameters!");
        println!("   This model cannot be trained.");
    } else if has_nan || has_inf {
        println!("⚠️  PARTIALLY WORKING: Model has parameters but produces NaN/Inf");
        println!("   Check initialization or numerical stability.");
    } else if !has_nonzero {
        println!("⚠️  PARTIALLY WORKING: Model produces only zeros");
        println!(
            "   Parameters exist ({}) but forward pass may be broken.",
            total_params
        );
    } else if step_result.is_err() {
        println!("⚠️  PARTIALLY WORKING: Forward pass works but optimizer fails");
        println!("   Error: {:?}", step_result.err());
    } else {
        println!(
            "✅ FUNCTIONAL: Model has {} parameters and produces valid outputs",
            total_params
        );
        println!("   Forward pass works correctly.");
        println!("   Optimizer integration successful.");
        println!("\n   FP32 TFT model is ready for training!");
    }

    Ok(())
}