Initial commit

crates/training/rtx-transformers/src/layers/mamba_tests.rs

@@ -0,0 +1,340 @@
+//! Comprehensive tests for Mamba (State Space Model) implementation
+//!
+//! This test suite covers:
+//! - MambaConfig validation and creation  
+//! - SelectiveScan operation correctness
+//! - State space discretization
+//! - Linear complexity attention alternative
+//! - MambaBlock forward/backward passes
+//! - State caching for inference
+//! - Integration with transformer infrastructure
+
+use super::mamba::{MambaConfig, MambaBlock, SelectiveScan, StateCache, MambaOutput};
+use crate::Result;
+use rtx_tensor::{Tensor, Device};
+use std::collections::HashMap;
+
+/// Test MambaConfig creation and validation
+#[tokio::test]
+async fn test_mamba_config_creation() -> Result<()> {
+    let config = MambaConfig {
+        d_model: 768,
+        d_state: 16,
+        d_conv: 4,
+        expand: 2,
+        dt_rank: None, // Auto-calculated
+        d_inner: None, // Auto-calculated
+        dt_min: 0.001,
+        dt_max: 0.1,
+        dt_init: "random".to_string(),
+        dt_scale: 1.0,
+        dt_init_floor: 1e-4,
+        conv_bias: true,
+        bias: false,
+        use_fast_path: true,
+    };
+
+    // Verify auto-calculated values
+    assert_eq!(config.get_d_inner(), 768 * 2); // d_model * expand
+    assert_eq!(config.get_dt_rank(), (768 + 15) / 16); // (d_model + 15) // 16
+    
+    Ok(())
+}
+
+/// Test basic selective scan operation
+#[tokio::test] 
+async fn test_selective_scan_basic() -> Result<()> {
+    let device = Device::cpu();
+    let batch_size = 2;
+    let seq_len = 10;
+    let d_inner = 16;
+    let d_state = 8;
+    
+    // Create input tensors
+    let u = Tensor::randn(&[batch_size, seq_len, d_inner], &device)?; // Input
+    let delta = Tensor::randn(&[batch_size, seq_len, d_inner], &device)?; // Time step
+    let A = Tensor::randn(&[d_inner, d_state], &device)?; // State transition
+    let B = Tensor::randn(&[batch_size, seq_len, d_state], &device)?; // Input matrix
+    let C = Tensor::randn(&[batch_size, seq_len, d_state], &device)?; // Output matrix
+    let D = Some(Tensor::randn(&[d_inner], &device)?); // Skip connection
+    
+    let selective_scan = SelectiveScan::new(d_inner, d_state);
+    let output = selective_scan.forward(&u, &delta, &A, &B, &C, D.as_ref())?;
+    
+    // Verify output shape
+    assert_eq!(output.shape().dims(), &[batch_size, seq_len, d_inner]);
+    
+    // Verify linear complexity by checking computation doesn't involve quadratic terms
+    // This is implicit in the state space formulation
+    Ok(())
+}
+
+/// Test state discretization correctness
+#[tokio::test]
+async fn test_state_discretization() -> Result<()> {
+    let device = Device::cpu();
+    let d_inner = 16;
+    let d_state = 8;
+    let seq_len = 5;
+    
+    let A = Tensor::randn(&[d_inner, d_state], &device)?;
+    let delta = Tensor::randn(&[1, seq_len, d_inner], &device)?;
+    
+    let selective_scan = SelectiveScan::new(d_inner, d_state);
+    let (deltaA, deltaB_u) = selective_scan.discretize(&delta, &A)?;
+    
+    // Verify discretized matrices have correct shapes
+    assert_eq!(deltaA.shape().dims(), &[1, seq_len, d_inner, d_state]);
+    
+    // deltaB_u should preserve sequence information
+    assert_eq!(deltaB_u.shape().dims()[1], seq_len);
+    
+    Ok(())
+}
+
+/// Test selective scan with state caching
+#[tokio::test]
+async fn test_selective_scan_with_cache() -> Result<()> {
+    let device = Device::cpu();
+    let config = MambaConfig::new(256, 16, 4);
+    let selective_scan = SelectiveScan::new(config.get_d_inner(), config.d_state);
+    
+    let batch_size = 1;
+    let seq_len_1 = 5;
+    let seq_len_2 = 3;
+    
+    // First sequence
+    let u1 = Tensor::randn(&[batch_size, seq_len_1, config.get_d_inner()], &device)?;
+    let delta1 = Tensor::randn(&[batch_size, seq_len_1, config.get_d_inner()], &device)?;
+    let A = Tensor::randn(&[config.get_d_inner(), config.d_state], &device)?;
+    let B1 = Tensor::randn(&[batch_size, seq_len_1, config.d_state], &device)?;
+    let C1 = Tensor::randn(&[batch_size, seq_len_1, config.d_state], &device)?;
+    
+    let mut cache = StateCache::new(batch_size, config.d_state, &device)?;
+    let output1 = selective_scan.forward_with_cache(&u1, &delta1, &A, &B1, &C1, None, &mut cache)?;
+    
+    // Second sequence (continues from first)
+    let u2 = Tensor::randn(&[batch_size, seq_len_2, config.get_d_inner()], &device)?;
+    let delta2 = Tensor::randn(&[batch_size, seq_len_2, config.get_d_inner()], &device)?;
+    let B2 = Tensor::randn(&[batch_size, seq_len_2, config.d_state], &device)?;
+    let C2 = Tensor::randn(&[batch_size, seq_len_2, config.d_state], &device)?;
+    
+    let output2 = selective_scan.forward_with_cache(&u2, &delta2, &A, &B2, &C2, None, &mut cache)?;
+    
+    // Verify outputs have correct shapes
+    assert_eq!(output1.shape().dims(), &[batch_size, seq_len_1, config.get_d_inner()]);
+    assert_eq!(output2.shape().dims(), &[batch_size, seq_len_2, config.get_d_inner()]);
+    
+    // Cache should maintain state between calls
+    assert!(!cache.is_empty());
+    
+    Ok(())
+}
+
+/// Test MambaBlock forward pass
+#[tokio::test]
+async fn test_mamba_block_forward() -> Result<()> {
+    let device = Device::cpu();
+    let config = MambaConfig::new(512, 16, 4);
+    let mamba_block = MambaBlock::new(config.clone(), &device)?;
+    
+    let batch_size = 2;
+    let seq_len = 20;
+    let input = Tensor::randn(&[batch_size, seq_len, config.d_model], &device)?;
+    
+    let output = mamba_block.forward(&input)?;
+    
+    // Output should have same shape as input
+    assert_eq!(output.output.shape().dims(), &[batch_size, seq_len, config.d_model]);
+    
+    // Verify output contains actual data (not just zeros/nans)
+    let output_data = output.output.to_vec()?;
+    assert!(!output_data.iter().all(|&x| x == 0.0));
+    assert!(output_data.iter().all(|&x| x.is_finite()));
+    
+    Ok(())
+}
+
+/// Test MambaBlock with residual connection
+#[tokio::test]
+async fn test_mamba_block_residual() -> Result<()> {
+    let device = Device::cpu();
+    let config = MambaConfig::new(256, 16, 4);
+    let mamba_block = MambaBlock::new(config.clone(), &device)?;
+    
+    let batch_size = 1;
+    let seq_len = 10;
+    let input = Tensor::randn(&[batch_size, seq_len, config.d_model], &device)?;
+    
+    let output = mamba_block.forward(&input)?;
+    
+    // With residual connections, output should differ from input
+    let input_norm = input.norm()?.to_scalar::<f32>()?;
+    let output_norm = output.output.norm()?.to_scalar::<f32>()?;
+    
+    // They should be different but in similar magnitude
+    assert!((input_norm - output_norm).abs() < input_norm * 2.0);
+    
+    Ok(())
+}
+
+/// Test gradient flow through MambaBlock
+#[tokio::test]
+async fn test_mamba_block_gradients() -> Result<()> {
+    let device = Device::cpu();
+    let config = MambaConfig::new(128, 8, 4);
+    let mut mamba_block = MambaBlock::new(config.clone(), &device)?;
+    
+    let batch_size = 1;
+    let seq_len = 5;
+    let input = Tensor::randn(&[batch_size, seq_len, config.d_model], &device)?;
+    input.set_requires_grad(true);
+    
+    let output = mamba_block.forward(&input)?;
+    let loss = output.output.mean()?;
+    
+    // Verify parameters can receive gradients
+    let params = mamba_block.parameters();
+    assert!(!params.is_empty());
+    
+    // Each parameter should have the correct shape
+    for param in params {
+        assert!(param.shape().numel() > 0);
+    }
+    
+    Ok(())
+}
+
+/// Test MambaBlock layer integration trait
+#[tokio::test]  
+async fn test_mamba_layer_trait() -> Result<()> {
+    let device = Device::cpu();
+    let config = MambaConfig::new(256, 16, 4);
+    let mamba_block = MambaBlock::new(config.clone(), &device)?;
+    
+    // Test Layer trait methods
+    assert_eq!(mamba_block.layer_type(), "MambaBlock");
+    assert_eq!(mamba_block.device(), &device);
+    
+    let params = mamba_block.parameters();
+    assert!(!params.is_empty());
+    
+    Ok(())
+}
+
+/// Test long sequence efficiency (linear complexity)
+#[tokio::test]
+async fn test_long_sequence_efficiency() -> Result<()> {
+    let device = Device::cpu();
+    let config = MambaConfig::new(128, 8, 4);
+    let mamba_block = MambaBlock::new(config.clone(), &device)?;
+    
+    // Test with increasingly longer sequences
+    let batch_size = 1;
+    let sequences = vec![10, 50, 100];
+    let mut processing_times = Vec::new();
+    
+    for seq_len in sequences {
+        let input = Tensor::randn(&[batch_size, seq_len, config.d_model], &device)?;
+        
+        let start = std::time::Instant::now();
+        let _output = mamba_block.forward(&input)?;
+        let elapsed = start.elapsed();
+        
+        processing_times.push(elapsed.as_nanos() as f64);
+        
+        // Verify output shape
+        assert_eq!(_output.output.shape().dims(), &[batch_size, seq_len, config.d_model]);
+    }
+    
+    // Should scale roughly linearly (not quadratically like attention)
+    // This is a basic check - exact scaling depends on implementation details
+    let ratio_1_2 = processing_times[1] / processing_times[0];
+    let ratio_2_3 = processing_times[2] / processing_times[1];
+    
+    // Ratios should be reasonable (not exponential growth)
+    assert!(ratio_1_2 < 10.0); // Generous bounds for test stability
+    assert!(ratio_2_3 < 5.0);
+    
+    Ok(())
+}
+
+/// Test conv1d operation in MambaBlock
+#[tokio::test]
+async fn test_mamba_conv1d() -> Result<()> {
+    let device = Device::cpu();
+    let config = MambaConfig::new(256, 16, 4);
+    let mamba_block = MambaBlock::new(config.clone(), &device)?;
+    
+    let batch_size = 2;
+    let seq_len = 15;
+    let input = Tensor::randn(&[batch_size, seq_len, config.d_model], &device)?;
+    
+    let output = mamba_block.forward(&input)?;
+    
+    // Conv1d should preserve sequence relationships
+    assert_eq!(output.output.shape().dims(), input.shape().dims());
+    
+    // Output should be different from input (conv1d transforms the data)
+    let input_sum = input.sum()?.to_scalar::<f32>()?;
+    let output_sum = output.output.sum()?.to_scalar::<f32>()?;
+    
+    assert!((input_sum - output_sum).abs() > 1e-6);
+    
+    Ok(())
+}
+
+/// Test state space parameters initialization
+#[tokio::test]
+async fn test_state_space_initialization() -> Result<()> {
+    let device = Device::cpu();
+    let config = MambaConfig::new(128, 8, 4);
+    let mamba_block = MambaBlock::new(config.clone(), &device)?;
+    
+    // Access internal state space parameters (if exposed)
+    let params = mamba_block.parameters();
+    
+    // Should have multiple parameter tensors for A, B, C matrices, etc.
+    assert!(params.len() >= 4); // At minimum: A, B, C, and D parameters
+    
+    // Each parameter should be properly initialized (not zero or NaN)
+    for param in params {
+        let param_data = param.to_vec()?;
+        assert!(!param_data.is_empty());
+        assert!(param_data.iter().all(|&x| x.is_finite()));
+        // Parameters shouldn't all be zero
+        assert!(param_data.iter().any(|&x| x != 0.0));
+    }
+    
+    Ok(())
+}
+
+/// Integration test with multiple MambaBlocks (stack)
+#[tokio::test]
+async fn test_mamba_stack() -> Result<()> {
+    let device = Device::cpu();
+    let config = MambaConfig::new(256, 16, 4);
+    
+    // Create a stack of 3 MambaBlocks  
+    let block1 = MambaBlock::new(config.clone(), &device)?;
+    let block2 = MambaBlock::new(config.clone(), &device)?;
+    let block3 = MambaBlock::new(config.clone(), &device)?;
+    
+    let batch_size = 1;
+    let seq_len = 12;
+    let mut x = Tensor::randn(&[batch_size, seq_len, config.d_model], &device)?;
+    
+    // Forward through stack
+    x = block1.forward(&x)?.output;
+    x = block2.forward(&x)?.output; 
+    x = block3.forward(&x)?.output;
+    
+    // Final output should have correct shape
+    assert_eq!(x.shape().dims(), &[batch_size, seq_len, config.d_model]);
+    
+    // Should produce meaningful transformations
+    let final_norm = x.norm()?.to_scalar::<f32>()?;
+    assert!(final_norm > 0.0 && final_norm.is_finite());
+    
+    Ok(())
+}