foxhunt/crates/ml/tests/liquid_cfc_training_test.rs

#![allow(
    clippy::assertions_on_constants,
    clippy::assertions_on_result_states,
    clippy::clone_on_copy,
    clippy::decimal_literal_representation,
    clippy::doc_markdown,
    clippy::empty_line_after_doc_comments,
    clippy::field_reassign_with_default,
    clippy::get_unwrap,
    clippy::identity_op,
    clippy::inconsistent_digit_grouping,
    clippy::indexing_slicing,
    clippy::integer_division,
    clippy::len_zero,
    clippy::let_underscore_must_use,
    clippy::manual_div_ceil,
    clippy::manual_let_else,
    clippy::manual_range_contains,
    clippy::modulo_arithmetic,
    clippy::needless_range_loop,
    clippy::non_ascii_literal,
    clippy::redundant_clone,
    clippy::shadow_reuse,
    clippy::shadow_same,
    clippy::shadow_unrelated,
    clippy::single_match_else,
    clippy::str_to_string,
    clippy::string_slice,
    clippy::tests_outside_test_module,
    clippy::too_many_lines,
    clippy::unnecessary_wraps,
    clippy::unseparated_literal_suffix,
    clippy::use_debug,
    clippy::useless_vec,
    clippy::wildcard_enum_match_arm,
    clippy::else_if_without_else,
    clippy::expect_used,
    clippy::missing_const_for_fn,
    clippy::similar_names,
    clippy::type_complexity,
    clippy::collapsible_else_if,
    clippy::doc_lazy_continuation,
    clippy::items_after_test_module,
    clippy::map_clone,
    clippy::multiple_unsafe_ops_per_block,
    clippy::unwrap_or_default,
    clippy::assign_op_pattern,
    clippy::needless_borrow,
    clippy::println_empty_string,
    clippy::unnecessary_cast,
    clippy::used_underscore_binding,
    clippy::create_dir,
    clippy::implicit_saturating_sub,
    clippy::exit,
    clippy::expect_fun_call,
    clippy::too_many_arguments,
    clippy::unnecessary_map_or,
    clippy::unwrap_used,
    dead_code,
    unused_imports,
    unused_variables,
    clippy::cloned_ref_to_slice_refs,
    clippy::neg_multiply,
    clippy::while_let_loop,
    clippy::bool_assert_comparison,
    clippy::excessive_precision,
    clippy::trivially_copy_pass_by_ref,
    clippy::op_ref,
    clippy::redundant_closure,
    clippy::unnecessary_lazy_evaluations,
    clippy::if_then_some_else_none,
    clippy::unnecessary_to_owned,
    clippy::single_component_path_imports,
)]
//! Integration test for Liquid CfC v2 full training loop
//!
//! Uses GPU-native GpuTensor through the LiquidTrainableAdapter's
//! public forward(&GpuTensor) and compute_loss(&GpuTensor, &GpuTensor) API.

use ml::liquid::adapter::LiquidTrainableAdapter;
use ml::liquid::candle_cfc::{CfCTrainConfig, DeviceConfig};
use ml::training::unified_trainer::UnifiedTrainable;
use ml_core::cuda_autograd::GpuTensor;

#[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::Cuda(0),
        learning_rate: 0.01,
        ..CfCTrainConfig::default()
    };

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

    // Synthetic training data
    let mut losses = Vec::new();
    for _ in 0..20 {
        // input: [batch=4, input_size=8] (Liquid adapter's forward expects 2D)
        let input = GpuTensor::randn(&[4, 8], 0.5, &stream).unwrap();
        let target = GpuTensor::zeros(&[4, 3], &stream).unwrap();

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

        adapter.backward(&loss_val).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::Cuda(0),
        ..CfCTrainConfig::default()
    };

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

    // Train a bit
    for _ in 0..5 {
        let input = GpuTensor::randn(&[2, 4], 1.0, &stream).unwrap();
        let target = GpuTensor::zeros(&[2, 2], &stream).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 metadata file 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 on same input
    let test_input = GpuTensor::randn(&[1, 4], 1.0, &stream).unwrap();
    let out1 = adapter.forward(&test_input).unwrap();
    let out2 = adapter2.forward(&test_input).unwrap();

    let host1 = out1.to_host(&stream).unwrap();
    let host2 = out2.to_host(&stream).unwrap();
    let diff: f32 = host1
        .iter()
        .zip(host2.iter())
        .map(|(a, b)| (a - b).abs())
        .sum();
    assert!(
        diff < 1e-3,
        "Checkpoint roundtrip should produce similar outputs, diff={}",
        diff
    );

    // Cleanup
    let _ = std::fs::remove_file(format!("{}.json", path_str));
    let _ = std::fs::remove_file(format!("{}.weights.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::Cuda(0),
        ..CfCTrainConfig::default()
    };

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

    // Validate by computing average loss over validation data
    let mut total_loss = 0.0;
    let val_count = 5;
    for _ in 0..val_count {
        let input = GpuTensor::randn(&[2, 4], 1.0, &stream).unwrap();
        let target = GpuTensor::zeros(&[2, 2], &stream).unwrap();
        let output = adapter.forward(&input).unwrap();
        let loss = adapter.compute_loss(&output, &target).unwrap();
        total_loss += loss;
    }

    let val_loss = total_loss / val_count as f64;
    assert!(val_loss.is_finite());
    assert!(val_loss >= 0.0);
}