foxhunt/ml/tests/xlstm_integration.rs

//! xLSTM (Extended Long Short-Term Memory) Integration Tests
//!
//! Validates the xLSTM trainable adapter end-to-end:
//! construction, forward pass, training loop, checkpoint save/load.

#![allow(unused_crate_dependencies)]

use candle_core::Device;
use candle_core::Tensor;
use ml::training::unified_trainer::UnifiedTrainable;
use ml::xlstm::config::XLSTMConfig;
use ml::xlstm::trainable::XLSTMTrainableAdapter;

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, &Device::Cpu);
    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_pass_3d() {
    let config = small_xlstm_config();
    let mut adapter = XLSTMTrainableAdapter::new(config, &Device::Cpu).unwrap();

    // [batch=4, seq_len=3, input_dim=8]
    let input = Tensor::randn(0f32, 1.0, (4, 3, 8), &Device::Cpu).unwrap();
    let output = adapter.forward(&input);
    assert!(output.is_ok(), "Forward failed: {:?}", output.err());

    let output = output.unwrap();
    // Output should be [batch, output_dim] = [4, 1]
    assert_eq!(output.dims(), &[4, 1], "Expected [4, 1] output shape");
}

#[test]
fn test_xlstm_forward_pass_2d() {
    let config = small_xlstm_config();
    let mut adapter = XLSTMTrainableAdapter::new(config, &Device::Cpu).unwrap();

    // Single-step input: [batch=4, input_dim=8]
    let input = Tensor::randn(0f32, 1.0, (4, 8), &Device::Cpu).unwrap();
    let output = adapter.forward(&input);
    assert!(output.is_ok(), "Forward (2D) failed: {:?}", output.err());

    let output = output.unwrap();
    assert_eq!(output.dims(), &[4, 1], "Expected [4, 1] output shape for 2D input");
}

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

    let batch_size = 8;
    let seq_len = 4;
    let input_dim = 8;
    let input =
        Tensor::randn(0f32, 0.5, (batch_size, seq_len, input_dim), &Device::Cpu).unwrap();

    // Target: simple learnable signal [batch, output_dim=1]
    let target = Tensor::randn(0f32, 0.1, (batch_size, 1), &Device::Cpu).unwrap();

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

    for epoch in 0..30 {
        let predictions = adapter.forward(&input).unwrap();

        let loss = adapter.compute_loss(&predictions, &target).unwrap();
        let loss_val = loss.to_scalar::<f32>().unwrap() as f64;

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

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

        if epoch % 10 == 0 {
            println!("xLSTM epoch {}: loss = {:.6}", epoch, loss_val);
        }
    }

    let first = first_loss.unwrap();
    println!(
        "xLSTM training: first_loss={:.6}, last_loss={:.6}",
        first, last_loss
    );

    // 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(), &Device::Cpu).unwrap();

    let input = Tensor::randn(0f32, 1.0, (4, 3, 8), &Device::Cpu).unwrap();
    let target = Tensor::randn(0f32, 0.1, (4, 1), &Device::Cpu).unwrap();

    // Train a few steps so weights diverge from initialization
    for _ in 0..3 {
        let pred = adapter.forward(&input).unwrap();
        let loss = adapter.compute_loss(&pred, &target).unwrap();
        adapter.backward(&loss).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, &Device::Cpu).unwrap();
    let load_result = adapter2.load_checkpoint(checkpoint_path.to_str().unwrap());
    assert!(
        load_result.is_ok(),
        "Load failed: {:?}",
        load_result.err()
    );

    // Compare predictions — should be identical after checkpoint restore
    let pred1 = adapter.forward(&input).unwrap();
    let pred2 = adapter2.forward(&input).unwrap();

    let diff = (pred1 - pred2)
        .unwrap()
        .abs()
        .unwrap()
        .sum_all()
        .unwrap()
        .to_scalar::<f32>()
        .unwrap();
    assert!(
        diff < 1e-4,
        "Checkpoint roundtrip predictions differ by {}",
        diff
    );

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

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

    // Create validation dataset with 3D inputs and matching [batch, 1] targets
    let val_data: Vec<(Tensor, Tensor)> = (0..5)
        .map(|_| {
            let input = Tensor::randn(0f32, 1.0, (4, 3, 8), &Device::Cpu).unwrap();
            let target = Tensor::randn(0f32, 0.1, (4, 1), &Device::Cpu).unwrap();
            (input, target)
        })
        .collect();

    let val_loss = adapter.validate(&val_data);
    assert!(
        val_loss.is_ok(),
        "Validation failed: {:?}",
        val_loss.err()
    );
    let loss_val = val_loss.unwrap();
    assert!(
        loss_val.is_finite(),
        "Validation loss is not finite: {}",
        loss_val
    );
    println!("xLSTM validation loss: {:.6}", loss_val);
}

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

    let input = Tensor::randn(0f32, 1.0, (4, 3, 8), &Device::Cpu).unwrap();
    let target = Tensor::randn(0f32, 0.1, (4, 1), &Device::Cpu).unwrap();

    let pred = adapter.forward(&input).unwrap();
    let loss = adapter.compute_loss(&pred, &target).unwrap();
    adapter.backward(&loss).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"));
}