Initial commit

crates/models/rtx-multimodal/src/vision/vit_tests.rs

@@ -0,0 +1,313 @@
+//! Comprehensive TDD tests for Vision Transformer (ViT)
+//! Following Red-Green-Refactor cycle with full implementations
+
+use super::*;
+use crate::MultimodalError;
+use rtx_tensor::{Device, Tensor};
+
+/// Test Suite for PatchEmbedding
+mod patch_embedding_tests {
+    use super::*;
+
+    #[test]
+    fn test_patch_embedding_dimensions() {
+        // RED: Test that patch embedding produces correct dimensions
+        let device = Device::default();
+        let patch_embed = PatchEmbedding::new(224, 16, 3, 768, &device).unwrap();
+
+        // Verify dimensions
+        assert_eq!(patch_embed.num_patches(), 196); // (224/16)^2
+        assert_eq!(patch_embed.embed_dim(), 768);
+    }
+
+    #[test]
+    fn test_patch_embedding_forward_shape() {
+        // GREEN: Test forward pass produces correct output shape
+        let device = Device::default();
+        let patch_embed = PatchEmbedding::new(32, 8, 3, 128, &device).unwrap();
+
+        // Create input tensor [batch=2, channels=3, height=32, width=32]
+        let input = Tensor::randn(&[2, 3, 32, 32], &device).unwrap();
+        let output = patch_embed.forward(&input).unwrap();
+
+        // Verify output shape [batch=2, num_patches=16, embed_dim=128]
+        assert_eq!(output.shape().dims(), &[2, 16, 128]);
+    }
+
+    #[test]
+    fn test_patch_embedding_invalid_size() {
+        // Test error handling for invalid patch size
+        let device = Device::default();
+        let result = PatchEmbedding::new(224, 15, 3, 768, &device);
+
+        assert!(result.is_err());
+        if let Err(e) = result {
+            match e {
+                MultimodalError::InvalidPatchSize {
+                    patch_size,
+                    image_size,
+                } => {
+                    assert_eq!(patch_size, 15);
+                    assert_eq!(image_size, 224);
+                }
+                _ => panic!("Expected InvalidPatchSize error"),
+            }
+        }
+    }
+
+    #[test]
+    fn test_patch_embedding_different_sizes() {
+        // Test various valid image and patch size combinations
+        let device = Device::default();
+        let test_cases = [
+            (64, 8, 8),   // 8x8 patches
+            (128, 16, 8), // 8x8 patches
+            (256, 32, 8), // 8x8 patches
+        ];
+
+        for (img_size, patch_size, expected_patches_per_dim) in test_cases {
+            let patch_embed = PatchEmbedding::new(img_size, patch_size, 3, 512, &device).unwrap();
+            let expected_total = expected_patches_per_dim * expected_patches_per_dim;
+            assert_eq!(patch_embed.num_patches(), expected_total);
+        }
+    }
+}
+
+/// Test Suite for PositionalEncoding
+mod positional_encoding_tests {
+    use super::*;
+
+    #[test]
+    fn test_positional_encoding_shape() {
+        let device = Device::default();
+        let pos_encoding = PositionalEncoding::new(197, 768, &device).unwrap();
+
+        // Test with different sequence lengths
+        let encoding = pos_encoding.forward(100).unwrap();
+        assert_eq!(encoding.shape().dims(), &[1, 100, 768]);
+    }
+
+    #[test]
+    fn test_positional_encoding_max_length() {
+        let device = Device::default();
+        let pos_encoding = PositionalEncoding::new(50, 256, &device).unwrap();
+
+        // Within max length - should succeed
+        let result = pos_encoding.forward(50);
+        assert!(result.is_ok());
+
+        // Exceeds max length - should fail
+        let result = pos_encoding.forward(51);
+        assert!(result.is_err());
+    }
+
+    #[test]
+    fn test_positional_encoding_consistency() {
+        let device = Device::default();
+        let pos_encoding = PositionalEncoding::new(100, 512, &device).unwrap();
+
+        // Same position should give same encoding
+        let encoding1 = pos_encoding.forward(50).unwrap();
+        let encoding2 = pos_encoding.forward(50).unwrap();
+
+        // Compare data (simplified - in practice would check actual values)
+        assert_eq!(encoding1.shape(), encoding2.shape());
+    }
+}
+
+/// Test Suite for MultiHeadAttention
+mod multihead_attention_tests {
+    use super::*;
+
+    #[test]
+    fn test_attention_creation() {
+        let device = Device::default();
+
+        // Valid configuration
+        let attention = MultiHeadAttention::new(768, 12, &device);
+        assert!(attention.is_ok());
+
+        // Invalid configuration (embed_dim not divisible by num_heads)
+        let attention = MultiHeadAttention::new(768, 13, &device);
+        assert!(attention.is_err());
+    }
+
+    #[test]
+    fn test_attention_forward_shape() {
+        let device = Device::default();
+        let attention = MultiHeadAttention::new(256, 8, &device).unwrap();
+
+        // Input: [batch=2, seq_len=10, embed_dim=256]
+        let input = Tensor::randn(&[2, 10, 256], &device).unwrap();
+        let output = attention.forward(&input).unwrap();
+
+        // Output should have same shape as input
+        assert_eq!(output.shape().dims(), &[2, 10, 256]);
+    }
+
+    #[test]
+    fn test_attention_invalid_input_dims() {
+        let device = Device::default();
+        let attention = MultiHeadAttention::new(256, 8, &device).unwrap();
+
+        // Invalid input shape (2D instead of 3D)
+        let input = Tensor::randn(&[10, 256], &device).unwrap();
+        let result = attention.forward(&input);
+
+        assert!(result.is_err());
+    }
+
+    #[test]
+    fn test_attention_scaling() {
+        let device = Device::default();
+        let embed_dim = 256;
+        let num_heads = 8;
+        let attention = MultiHeadAttention::new(embed_dim, num_heads, &device).unwrap();
+
+        // Scale factor is internal - we verify it through forward pass behavior
+    }
+}
+
+/// Test Suite for complete VisionTransformer
+mod vision_transformer_tests {
+    use super::*;
+
+    #[test]
+    fn test_vit_creation() {
+        let device = Device::default();
+        let config = VisionConfig::new(224, 16, 768, 12);
+
+        let vit = VisionTransformer::new(config.clone(), &device);
+        assert!(vit.is_ok());
+    }
+
+    #[test]
+    fn test_vit_forward_classification() {
+        let device = Device::default();
+        let config = VisionConfig::new(32, 8, 64, 4);
+
+        let mut vit = VisionTransformer::new(config.clone(), &device).unwrap();
+
+        // Input image [batch=2, channels=3, height=32, width=32]
+        let input = Tensor::randn(&[2, 3, 32, 32], &device).unwrap();
+        let output = vit.forward(&input).unwrap();
+
+        // Output should be [batch=2, seq_len=17, hidden_dim=64]
+        // seq_len = 1 (cls token) + 16 (4x4 patches from 32x32 image with patch_size=8)
+        let expected_seq_len = 1 + (32 / 8) * (32 / 8); // 1 + 16 = 17
+        assert_eq!(output.shape().dims(), &[2, expected_seq_len, 64]);
+    }
+
+    #[test]
+    fn test_vit_extract_features() {
+        let device = Device::default();
+        let config = VisionConfig::new(32, 8, 64, 4);
+
+        let mut vit = VisionTransformer::new(config.clone(), &device).unwrap();
+
+        let input = Tensor::randn(&[2, 3, 32, 32], &device).unwrap();
+        let features = vit.forward(&input).unwrap();
+
+        // Features should include cls token + patches
+        let expected_seq_len = 1 + (32 / 8) * (32 / 8); // 1 + 16 = 17
+        assert_eq!(features.shape().dims(), &[2, expected_seq_len, 64]);
+    }
+
+    #[test]
+    fn test_vit_batch_processing() {
+        let device = Device::default();
+        let config = VisionConfig::new(32, 8, 64, 4);
+
+        let mut vit = VisionTransformer::new(config.clone(), &device).unwrap();
+
+        // Test different batch sizes
+        let expected_seq_len = 1 + (32 / 8) * (32 / 8); // 1 + 16 = 17
+        for batch_size in [1, 2, 4, 8] {
+            let input = Tensor::randn(&[batch_size, 3, 32, 32], &device).unwrap();
+            let output = vit.forward(&input).unwrap();
+            assert_eq!(output.shape().dims()[0], batch_size);
+            assert_eq!(output.shape().dims()[1], expected_seq_len); // seq_len
+            assert_eq!(output.shape().dims()[2], 64); // hidden_dim
+        }
+    }
+}
+
+/// Integration tests for Vision Transformer components
+mod integration_tests {
+    use super::*;
+
+    #[test]
+    fn test_end_to_end_pipeline() {
+        let device = Device::default();
+
+        // Small configuration for testing
+        let config = VisionConfig::new(64, 16, 128, 4);
+
+        let mut vit = VisionTransformer::new(config.clone(), &device).unwrap();
+
+        // Process batch of images
+        let images = Tensor::randn(&[4, 3, 64, 64], &device).unwrap();
+        let output = vit.forward(&images).unwrap();
+
+        // Verify output shape: [batch, seq_len, hidden_dim]
+        let expected_seq_len = 1 + (64 / 16) * (64 / 16); // 1 + 16 = 17
+        assert_eq!(output.shape().dims(), &[4, expected_seq_len, 128]);
+
+        // Extract features for downstream tasks
+        let features = vit.forward(&images).unwrap();
+        assert_eq!(features.shape().dims()[0], 4); // batch size
+        assert_eq!(features.shape().dims()[1], expected_seq_len); // seq_len
+        assert_eq!(features.shape().dims()[2], 128); // embed_dim
+    }
+
+    #[test]
+    fn test_gradient_flow() {
+        let device = Device::default();
+        let config = VisionConfig::new(32, 8, 64, 4);
+
+        let mut vit = VisionTransformer::new(config.clone(), &device).unwrap();
+
+        // Forward pass
+        let input = Tensor::randn(&[2, 3, 32, 32], &device).unwrap();
+        let output = vit.forward(&input).unwrap();
+
+        // Verify gradients can flow (shape check as proxy)
+        let expected_seq_len = 1 + (32 / 8) * (32 / 8); // 1 + 16 = 17
+        assert_eq!(output.shape().dims(), &[2, expected_seq_len, 64]);
+
+        // In a real implementation, we would:
+        // 1. Compute loss
+        // 2. Backward pass
+        // 3. Check gradients exist and are non-zero
+    }
+}
+
+/// Performance and stress tests
+mod performance_tests {
+    use super::*;
+
+    #[test]
+    #[ignore] // Mark as ignored for regular test runs
+    fn test_large_model() {
+        let device = Device::default();
+        let config = VisionConfig::new(224, 16, 768, 12);
+
+        let vit = VisionTransformer::new(config.clone(), &device);
+        assert!(vit.is_ok());
+    }
+
+    #[test]
+    fn test_memory_efficiency() {
+        let device = Device::default();
+        let config = VisionConfig::new(32, 8, 64, 4);
+
+        let mut vit = VisionTransformer::new(config.clone(), &device).unwrap();
+
+        // Process multiple batches without memory leak
+        for _ in 0..10 {
+            let input = Tensor::randn(&[4, 3, 32, 32], &device).unwrap();
+            let _ = vit.forward(&input).unwrap();
+        }
+        // In practice, would monitor memory usage
+    }
+}