foxhunt/ml/tests/liquid_cfc_training_test.rs

//! Integration test for Liquid CfC v2 full training loop

use candle_core::{DType, Tensor};
use ml::liquid::adapter::LiquidTrainableAdapter;
use ml::liquid::candle_cfc::{CfCTrainConfig, DeviceConfig};
use ml::training::unified_trainer::UnifiedTrainable;

#[test]
fn test_liquid_cfc_full_training_loop() {
    let config = CfCTrainConfig {
        input_size: 8,
        hidden_size: 32,
        output_size: 3,
        backbone_hidden_sizes: vec![32],
        seq_len: 10,
        device: DeviceConfig::Cpu,
        learning_rate: 0.01,
        ..CfCTrainConfig::default()
    };

    let mut adapter = LiquidTrainableAdapter::new(config).unwrap();
    let device = adapter.device().clone();

    // Synthetic training data
    let mut losses = Vec::new();
    for _ in 0..20 {
        let input = Tensor::randn(0f32, 0.5, (4, 10, 8), &device).unwrap();
        let target = Tensor::zeros((4, 3), DType::F32, &device).unwrap();

        let output = adapter.forward(&input).unwrap();
        let loss = adapter.compute_loss(&output, &target).unwrap();
        let loss_val: f32 = loss.to_scalar().unwrap();
        losses.push(loss_val);

        adapter.backward(&loss).unwrap();
        adapter.optimizer_step().unwrap();
    }

    // Verify training happened
    assert_eq!(adapter.get_step(), 20);
    assert_eq!(adapter.model_type(), "Liquid-CfC");

    // Loss should generally decrease (allow some noise)
    let first_5_avg: f32 = losses.iter().take(5).sum::<f32>() / 5.0;
    let last_5_avg: f32 = losses.iter().rev().take(5).sum::<f32>() / 5.0;
    assert!(
        last_5_avg < first_5_avg * 1.5,
        "Loss should trend down: first_5={:.4}, last_5={:.4}",
        first_5_avg,
        last_5_avg
    );
}

#[test]
fn test_liquid_cfc_checkpoint_roundtrip() {
    let config = CfCTrainConfig {
        input_size: 4,
        hidden_size: 8,
        output_size: 2,
        backbone_hidden_sizes: vec![8],
        seq_len: 3,
        device: DeviceConfig::Cpu,
        ..CfCTrainConfig::default()
    };

    let mut adapter = LiquidTrainableAdapter::new(config.clone()).unwrap();
    let device = adapter.device().clone();

    // Train a bit
    for _ in 0..5 {
        let input = Tensor::randn(0f32, 1.0, (2, 3, 4), &device).unwrap();
        let target = Tensor::zeros((2, 2), DType::F32, &device).unwrap();
        let output = adapter.forward(&input).unwrap();
        let loss = adapter.compute_loss(&output, &target).unwrap();
        adapter.backward(&loss).unwrap();
        adapter.optimizer_step().unwrap();
    }

    // Save checkpoint
    let tmp_dir = std::env::temp_dir().join("liquid_cfc_integration_test");
    let _ = std::fs::create_dir_all(&tmp_dir);
    let checkpoint_path = tmp_dir.join("liquid_test");
    let path_str = checkpoint_path.to_str().unwrap();
    adapter.save_checkpoint(path_str).unwrap();

    // Verify files exist
    assert!(std::path::Path::new(&format!("{}.safetensors", path_str)).exists());
    assert!(std::path::Path::new(&format!("{}.json", path_str)).exists());

    // Load into new adapter
    let mut adapter2 = LiquidTrainableAdapter::new(config).unwrap();
    let metadata = adapter2.load_checkpoint(path_str).unwrap();
    assert_eq!(metadata.model_type, "Liquid-CfC");
    assert_eq!(metadata.step, 5);

    // Verify same predictions
    let test_input = Tensor::randn(0f32, 1.0, (1, 3, 4), &device).unwrap();
    let out1 = adapter.forward(&test_input).unwrap();
    let out2 = adapter2.forward(&test_input).unwrap();

    let diff: f32 = (&out1 - &out2)
        .unwrap()
        .abs()
        .unwrap()
        .sum_all()
        .unwrap()
        .to_scalar()
        .unwrap();
    assert!(
        diff < 1e-5,
        "Checkpoint roundtrip should produce identical outputs, diff={}",
        diff
    );

    // Cleanup
    let _ = std::fs::remove_file(format!("{}.safetensors", path_str));
    let _ = std::fs::remove_file(format!("{}.json", path_str));
    let _ = std::fs::remove_dir(&tmp_dir);
}

#[test]
fn test_liquid_cfc_validate() {
    let config = CfCTrainConfig {
        input_size: 4,
        hidden_size: 8,
        output_size: 2,
        backbone_hidden_sizes: vec![8],
        seq_len: 3,
        device: DeviceConfig::Cpu,
        ..CfCTrainConfig::default()
    };

    let mut adapter = LiquidTrainableAdapter::new(config).unwrap();
    let device = adapter.device().clone();

    let val_data: Vec<(Tensor, Tensor)> = (0..5)
        .map(|_| {
            (
                Tensor::randn(0f32, 1.0, (2, 3, 4), &device).unwrap(),
                Tensor::zeros((2, 2), DType::F32, &device).unwrap(),
            )
        })
        .collect();

    let val_loss = adapter.validate(&val_data).unwrap();
    assert!(val_loss.is_finite());
    assert!(val_loss >= 0.0);
}