rustytorch/crates/training/rtx-transformers/simple_alibi_test.rs

// Simple standalone ALiBi test that doesn't depend on external crates
// Run with: rustc simple_alibi_test.rs && ./simple_alibi_test

use std::collections::HashMap;

// Basic mock types for testing ALiBi functionality
#[derive(Debug, Clone, PartialEq)]
pub enum Device {
    Cpu,
}

#[derive(Debug, Clone, Copy)]
pub enum DType {
    F32,
}

#[derive(Debug, Clone)]
pub struct Tensor {
    shape: Vec<usize>,
}

impl Tensor {
    pub fn zeros(shape: Vec<usize>, _dtype: DType, _device: &Device) -> std::result::Result<Self, String> {
        Ok(Self { shape })
    }
    
    pub fn ones(shape: Vec<usize>, _dtype: DType, _device: &Device) -> std::result::Result<Self, String> {
        Ok(Self { shape })
    }
    
    pub fn shape(&self) -> &[usize] {
        &self.shape
    }
}

#[derive(Debug)]
pub enum TransformerError {
    Config(String),
    TensorOp(String),
}

impl TransformerError {
    pub fn config(message: String) -> Self {
        Self::Config(message)
    }
    
    pub fn tensor_op(message: String) -> Self {
        Self::TensorOp(message)
    }
}

impl std::fmt::Display for TransformerError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            TransformerError::Config(msg) => write!(f, "Config error: {}", msg),
            TransformerError::TensorOp(msg) => write!(f, "Tensor error: {}", msg),
        }
    }
}

impl std::error::Error for TransformerError {}

pub type Result<T> = std::result::Result<T, TransformerError>;

// ALiBi Implementation (same as in our actual file)
#[derive(Debug, Clone)]
pub struct AliBiConfig {
    pub num_heads: usize,
    pub causal: bool,
    pub max_seq_len: usize,
    pub device: Device,
}

impl Default for AliBiConfig {
    fn default() -> Self {
        Self {
            num_heads: 8,
            causal: true,
            max_seq_len: 2048,
            device: Device::cuda(0).unwrap_or(Device::default()),
        }
    }
}

#[derive(Debug)]
struct BiasCache {
    cache: std::sync::RwLock<HashMap<(usize, bool), Tensor>>,
}

impl BiasCache {
    fn new() -> Self {
        Self {
            cache: std::sync::RwLock::new(HashMap::new()),
        }
    }

    fn get_or_compute(&self, seq_len: usize, causal: bool, slopes: &[f32], device: &Device) -> Result<Tensor> {
        let key = (seq_len, causal);
        
        {
            let cache = self.cache.read().unwrap();
            if let Some(bias) = cache.get(&key) {
                return Ok(bias.clone());
            }
        }
        
        let bias = Self::compute_bias_matrix(seq_len, causal, slopes, device)?;
        
        {
            let mut cache = self.cache.write().unwrap();
            cache.insert(key, bias.clone());
        }
        
        Ok(bias)
    }

    fn compute_bias_matrix(seq_len: usize, _causal: bool, slopes: &[f32], device: &Device) -> Result<Tensor> {
        let shape = vec![slopes.len(), seq_len, seq_len];
        Tensor::zeros(shape, DType::F32, device)
            .map_err(|e| TransformerError::tensor_op(format!("Failed to create bias matrix: {}", e)))
    }
}

#[derive(Debug)]
pub struct ALiBi {
    config: AliBiConfig,
    slopes: Vec<f32>,
    bias_cache: BiasCache,
}

impl ALiBi {
    pub fn new(config: AliBiConfig) -> Result<Self> {
        if config.num_heads == 0 {
            return Err(TransformerError::config("num_heads must be greater than 0".to_string()));
        }
        
        let slopes = Self::compute_slopes(config.num_heads);
        
        Ok(Self {
            config,
            slopes,
            bias_cache: BiasCache::new(),
        })
    }

    pub fn compute_slopes(num_heads: usize) -> Vec<f32> {
        if num_heads == 0 {
            return vec![];
        }
        
        (0..num_heads).map(|i| 2.0_f32.powf(-(8.0 * i as f32 / num_heads as f32))).collect()
    }

    pub fn get_bias(&self, seq_len: usize) -> Result<Tensor> {
        if seq_len == 0 {
            return Err(TransformerError::config("seq_len must be greater than 0".to_string()));
        }
        
        self.bias_cache.get_or_compute(seq_len, self.config.causal, &self.slopes, &self.config.device)
    }

    pub fn apply_bias(&self, attention_scores: &Tensor, seq_len: usize) -> Result<Tensor> {
        let _bias = self.get_bias(seq_len)?;
        Ok(attention_scores.clone())
    }

    pub fn get_slopes(&self) -> &[f32] {
        &self.slopes
    }

    pub fn config(&self) -> &AliBiConfig {
        &self.config
    }

    pub fn is_causal(&self) -> bool {
        self.config.causal
    }

    pub fn max_seq_len(&self) -> usize {
        self.config.max_seq_len
    }

    pub fn clear_cache(&self) {
        let mut cache = self.bias_cache.cache.write().unwrap();
        cache.clear();
    }

    pub fn cache_size(&self) -> usize {
        let cache = self.bias_cache.cache.read().unwrap();
        cache.len()
    }
}

fn run_test(test_name: &str, test_fn: fn() -> std::result::Result<(), Box<dyn std::error::Error>>) {
    print!("Running {}: ", test_name);
    match test_fn() {
        Ok(()) => println!("PASS ✓"),
        Err(e) => println!("FAIL ✗ - {}", e),
    }
}

fn main() -> std::result::Result<(), Box<dyn std::error::Error>> {
    println!("🧪 ALiBi Red Phase Tests (All should currently pass since we have placeholder implementations)");
    println!("Note: This demonstrates the TDD RED phase - tests exist but implementation is minimal");
    println!();

    run_test("Default config test", || {
        let config = AliBiConfig::default();
        assert_eq!(config.num_heads, 8);
        assert_eq!(config.causal, true);
        assert_eq!(config.max_seq_len, 2048);
        assert_eq!(config.device, Device::cuda(0).unwrap_or(Device::default()));
        Ok(())
    });

    run_test("ALiBi creation with valid config", || {
        let config = AliBiConfig {
            num_heads: 12,
            causal: false,
            max_seq_len: 4096,
            device: Device::cuda(0).unwrap_or(Device::default()),
        };
        
        let alibi = ALiBi::new(config)?;
        assert_eq!(alibi.config().num_heads, 12);
        assert_eq!(alibi.is_causal(), false);
        assert_eq!(alibi.max_seq_len(), 4096);
        assert_eq!(alibi.get_slopes().len(), 12);
        Ok(())
    });

    run_test("ALiBi creation with zero heads should fail", || {
        let config = AliBiConfig {
            num_heads: 0,
            causal: true,
            max_seq_len: 1024,
            device: Device::cuda(0).unwrap_or(Device::default()),
        };
        
        let result = ALiBi::new(config);
        assert!(result.is_err());
        let error_msg = result.unwrap_err().to_string();
        assert!(error_msg.contains("num_heads must be greater than 0"));
        Ok(())
    });

    run_test("Compute slopes for zero heads", || {
        let slopes = ALiBi::compute_slopes(0);
        assert!(slopes.is_empty());
        Ok(())
    });

    run_test("Compute slopes correct length", || {
        let slopes = ALiBi::compute_slopes(8);
        assert_eq!(slopes.len(), 8);
        Ok(())
    });

    run_test("Slopes are decreasing geometric sequence", || {
        let slopes = ALiBi::compute_slopes(4);
        assert_eq!(slopes.len(), 4);
        for i in 1..slopes.len() {
            if slopes[i] >= slopes[i-1] {
                return Err(format!("Slopes should be decreasing: {} >= {}", slopes[i], slopes[i-1]).as_str().into());
            }
        }
        Ok(())
    });

    run_test("Get bias with zero seq_len should fail", || {
        let config = AliBiConfig::default();
        let alibi = ALiBi::new(config)?;
        
        let result = alibi.get_bias(0);
        assert!(result.is_err());
        let error_msg = result.unwrap_err().to_string();
        assert!(error_msg.contains("seq_len must be greater than 0"));
        Ok(())
    });

    run_test("Get bias returns tensor with correct shape", || {
        let config = AliBiConfig::default();
        let alibi = ALiBi::new(config)?;
        
        let bias = alibi.get_bias(10)?;
        let shape = bias.shape();
        assert_eq!(shape.len(), 3);
        assert_eq!(shape[0], 8); // num_heads
        assert_eq!(shape[1], 10); // seq_len
        assert_eq!(shape[2], 10); // seq_len
        Ok(())
    });

    run_test("Apply bias preserves shape", || {
        let config = AliBiConfig::default();
        let alibi = ALiBi::new(config)?;
        
        let attention_scores = Tensor::ones(vec![8, 10, 10], DType::F32, &Device::cuda(0).unwrap_or(Device::default()))?;
        let result = alibi.apply_bias(&attention_scores, 10)?;
        
        assert_eq!(result.shape(), attention_scores.shape());
        Ok(())
    });

    run_test("Cache operations work correctly", || {
        let config = AliBiConfig::default();
        let alibi = ALiBi::new(config)?;
        
        assert_eq!(alibi.cache_size(), 0);
        
        let _bias = alibi.get_bias(10)?;
        assert_eq!(alibi.cache_size(), 1);
        
        alibi.clear_cache();
        assert_eq!(alibi.cache_size(), 0);
        Ok(())
    });

    run_test("Cache reuse works properly", || {
        let config = AliBiConfig::default();
        let alibi = ALiBi::new(config)?;
        
        let bias1 = alibi.get_bias(10)?;
        let bias2 = alibi.get_bias(10)?;
        
        // Should only have one cached entry
        assert_eq!(alibi.cache_size(), 1);
        assert_eq!(bias1.shape(), bias2.shape());
        Ok(())
    });

    println!();
    println!("🎯 RED Phase Complete!");
    println!("All basic structural tests pass with minimal/placeholder implementations.");
    println!("Next step: GREEN phase - implement actual ALiBi functionality to pass comprehensive tests.");

    Ok(())
}