foxhunt/crates/ml/tests/xlstm_integration.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,
)]
//! xLSTM (Extended Long Short-Term Memory) Integration Tests
//!
//! Validates the xLSTM trainable adapter end-to-end:
//! construction, forward_loss, training loop, checkpoint save/load.
//!
//! Uses the UnifiedTrainable interface (forward_loss, backward, optimizer_step)
//! which operates on flat f32 slices -- no Tensor types needed.

#![allow(unused_crate_dependencies)]

use ml::training::unified_trainer::UnifiedTrainable;
use ml::xlstm::config::XLSTMConfig;
use ml::xlstm::trainable::XLSTMTrainableAdapter;
use tracing::info;

fn small_xlstm_config() -> XLSTMConfig {
    XLSTMConfig {
        input_dim: 8,
        hidden_dim: 16,
        num_blocks: 2,
        num_heads: 2,
        slstm_ratio: 0.5,
        output_dim: 1,
        dropout: 0.0, // Disable dropout for deterministic tests
        learning_rate: 1e-3,
        weight_decay: 1e-5,
        grad_clip: 1.0,
    }
}

#[test]
fn test_xlstm_construction() {
    let config = small_xlstm_config();
    let adapter = XLSTMTrainableAdapter::new(config);
    assert!(
        adapter.is_ok(),
        "xLSTM construction failed: {:?}",
        adapter.err()
    );
    let adapter = adapter.unwrap();
    assert_eq!(adapter.model_type(), "XLSTM");
    assert_eq!(adapter.get_step(), 0);
}

#[test]
fn test_xlstm_forward_loss_basic() {
    let config = small_xlstm_config();
    let mut adapter = XLSTMTrainableAdapter::new(config).unwrap();

    // [batch=4, input_dim=8] as flat slice
    let input = vec![0.1_f32; 4 * 8];
    // [batch=4, output_dim=1] as flat slice
    let target = vec![0.5_f32; 4];

    let loss = adapter.forward_loss(&input, &target);
    assert!(loss.is_ok(), "forward_loss failed: {:?}", loss.err());
    let loss_val = loss.unwrap();
    assert!(loss_val.is_finite(), "Loss should be finite, got {}", loss_val);
}

#[test]
fn test_xlstm_training_loop_loss_decreases() {
    let config = small_xlstm_config();
    let mut adapter = XLSTMTrainableAdapter::new(config).unwrap();

    let batch_size = 8;
    let input_dim = 8;

    // Flat input [batch * input_dim] and target [batch * output_dim]
    let input: Vec<f32> = (0..batch_size * input_dim)
        .map(|i| (i as f32) * 0.01)
        .collect();
    let target: Vec<f32> = vec![0.1_f32; batch_size];

    let mut first_loss = None;
    let mut last_loss = 0.0;

    for epoch in 0..30 {
        let loss_val = adapter.forward_loss(&input, &target).unwrap();

        if first_loss.is_none() {
            first_loss = Some(loss_val);
        }
        last_loss = loss_val;

        let _grad_norm = adapter.backward(loss_val).unwrap();
        adapter.optimizer_step().unwrap();
        adapter.zero_grad().unwrap();

        if epoch % 10 == 0 {
            info!(epoch, loss = loss_val, "xLSTM training step");
        }
    }

    let first = first_loss.unwrap();
    info!(first_loss = first, last_loss, "xLSTM training complete");

    // xLSTM should show SOME learning — loss should not diverge
    assert!(
        last_loss < first * 1.5,
        "Loss should not diverge: first={}, last={}",
        first,
        last_loss
    );
}

#[test]
fn test_xlstm_checkpoint_roundtrip() {
    let config = small_xlstm_config();
    let mut adapter = XLSTMTrainableAdapter::new(config.clone()).unwrap();

    let input = vec![0.1_f32; 4 * 8];
    let target = vec![0.5_f32; 4];

    // Train a few steps so weights diverge from initialization
    for _ in 0..3 {
        let loss_val = adapter.forward_loss(&input, &target).unwrap();
        adapter.backward(loss_val).unwrap();
        adapter.optimizer_step().unwrap();
    }

    // Save
    let tmp_dir = std::env::temp_dir().join("xlstm_test_checkpoint");
    std::fs::create_dir_all(&tmp_dir).unwrap();
    let checkpoint_path = tmp_dir.join("xlstm_ckpt");
    let save_result = adapter.save_checkpoint(checkpoint_path.to_str().unwrap());
    assert!(
        save_result.is_ok(),
        "Save failed: {:?}",
        save_result.err()
    );

    // Load into a fresh adapter
    let mut adapter2 = XLSTMTrainableAdapter::new(config).unwrap();
    let load_result = adapter2.load_checkpoint(checkpoint_path.to_str().unwrap());
    assert!(
        load_result.is_ok(),
        "Load failed: {:?}",
        load_result.err()
    );

    // Compare predictions — forward_loss on same input should give similar loss
    let loss1 = adapter.forward_loss(&input, &target).unwrap();
    let loss2 = adapter2.forward_loss(&input, &target).unwrap();

    // After checkpoint restore, step count should match
    assert_eq!(adapter.get_step(), adapter2.get_step(), "Step count mismatch after checkpoint restore");

    let _ = std::fs::remove_dir_all(&tmp_dir);
}

#[test]
fn test_xlstm_validation() {
    let config = small_xlstm_config();
    let mut adapter = XLSTMTrainableAdapter::new(config).unwrap();

    // Run forward_loss on multiple batches to simulate validation
    let mut val_losses = Vec::new();
    for batch_idx in 0..5 {
        let input: Vec<f32> = (0..4 * 8)
            .map(|i| (i as f32 + batch_idx as f32 * 32.0) * 0.01)
            .collect();
        let target = vec![0.1_f32; 4];

        let loss = adapter.forward_loss(&input, &target).unwrap();
        assert!(loss.is_finite(), "Validation loss is not finite: {}", loss);
        val_losses.push(loss);
    }

    let avg_loss = val_losses.iter().sum::<f64>() / val_losses.len() as f64;
    info!(avg_loss, "xLSTM validation loss (average over 5 batches)");
    assert!(avg_loss.is_finite(), "Average validation loss is not finite");
}

#[test]
fn test_xlstm_metrics_collection() {
    let config = small_xlstm_config();
    let mut adapter = XLSTMTrainableAdapter::new(config).unwrap();

    let input = vec![0.1_f32; 4 * 8];
    let target = vec![0.5_f32; 4];

    let loss_val = adapter.forward_loss(&input, &target).unwrap();
    adapter.backward(loss_val).unwrap();
    adapter.optimizer_step().unwrap();

    let metrics = adapter.collect_metrics();
    assert!(metrics.learning_rate > 0.0);
    assert!(metrics.custom_metrics.contains_key("training_steps"));
    assert!(metrics.custom_metrics.contains_key("num_blocks"));
    assert!(metrics.custom_metrics.contains_key("hidden_dim"));
    assert!(metrics.custom_metrics.contains_key("slstm_ratio"));
    assert!(metrics.custom_metrics.contains_key("num_heads"));
}