foxhunt/AGENT_220_QUICK_REFERENCE.md

Agent 220: TDD Shape Tests - Quick Reference

For Developers: Fast guide to using the new MAMBA-2 shape tests.

---

🚀 Quick Start

Run All Tests

cd /home/jgrusewski/Work/foxhunt
cargo test -p ml --test mamba2_shape_tests -- --nocapture

Run Single Test

cargo test -p ml test_forward_pass_shapes -- --nocapture

Run Only Passing Tests (Fast Validation)

cargo test -p ml test_batch_concatenation -- --nocapture
cargo test -p ml test_optimizer_scalar_dtypes -- --nocapture
cargo test -p ml test_zero_sequence_length -- --nocapture

---

🐛 Current Status: Bug #18 Blocking Tests

Error: Model error: Layer normalization failed: unsupported dtype for rmsnorm F64

Impact: 11/18 tests blocked (all tests that call model.forward())

Fix Required: Convert F64 → F32 for LayerNorm operation

---

🔧 How to Fix Bug #18

File to Edit

/home/jgrusewski/Work/foxhunt/ml/src/mamba/mod.rs

Function to Update

CudaLayerNorm::forward

Recommended Fix

pub fn forward(&self, x: &Tensor) -> Result<Tensor, MLError> {
    // Convert F64 → F32 for LayerNorm (Candle only supports F32)
    let x_f32 = if x.dtype() == DType::F64 {
        x.to_dtype(DType::F32)?
    } else {
        x.clone()
    };

    // Apply LayerNorm
    let normalized = layer_norm_with_fallback(
        &x_f32,
        &self.normalized_shape,
        self.weight.as_ref(),
        self.bias.as_ref(),
        self.eps,
    )?;

    // Convert back to F64 to maintain model dtype
    if x.dtype() == DType::F64 {
        normalized.to_dtype(DType::F64)
            .map_err(|e| MLError::ModelError(format!("Failed to convert LayerNorm output to F64: {}", e)))
    } else {
        Ok(normalized)
    }
}

After Fixing

cargo test -p ml --test mamba2_shape_tests
# Expected: 18/18 tests pass

---

📊 Test Coverage Map

Test Name	Bugs Caught	Status	Duration
test_forward_pass_shapes	#1-5	🔴 Blocked	<1s
test_ssm_matrix_broadcast_shapes	#4	🔴 Blocked	<1s
test_loss_computation_shapes	#6	🔴 Blocked	<1s
test_all_tensors_dtype_f64	#7-10	🔴 Blocked	<1s
test_discretization_dtype_consistency	#8-9	🔴 Blocked	<1s
test_adam_optimizer_broadcasts	#11-14	🔴 Blocked	<1s
test_single_training_step	#15-17	🔴 Blocked	1-2s
test_validation_loss_consistency	#17	🔴 Blocked	<1s
test_full_training_cycle_integration	#1-17	🔴 Blocked	2-3s
test_single_sample_batch	Edge case	🔴 Blocked	<1s
test_large_batch_size	Stress test	🔴 Blocked	<1s
test_batch_concatenation	#15	✅ PASS	<1s
test_optimizer_scalar_dtypes	#12	✅ PASS	<1s
test_zero_sequence_length	Edge case	✅ PASS	<1s

Total Duration: 5-10 seconds (after Bug #18 fix)

---

🎯 What Each Test Validates

Shape Tests (Catch Bugs #1-6)

• test_forward_pass_shapes: Output is [batch, seq, d_model], not [batch, seq, 1]
• test_ssm_matrix_broadcast_shapes: B/C matrices broadcast across batch dimension
• test_loss_computation_shapes: Loss uses output_last [batch, 1, d_model]

Dtype Tests (Catch Bugs #7-10)

• test_all_tensors_dtype_f64: All SSM matrices are F64 (no F32 sneaks in)
• test_discretization_dtype_consistency: dt_mean uses F64 (not F32)
• test_optimizer_scalar_dtypes: Scalars match tensor dtype (F64 or F32)

Broadcast Tests (Catch Bugs #11-14)

• test_adam_optimizer_broadcasts: Adam scalars (beta1, beta2, lr) broadcast correctly
• test_optimizer_scalar_dtypes: Tensor::new uses correct dtype

Training Tests (Catch Bugs #15-17)

• test_single_training_step: End-to-end training (forward → loss → backward → optimize)
• test_batch_concatenation: Individual samples [1, seq, d_model] → batched [batch, seq, d_model]
• test_validation_loss_consistency: Validation uses output_last (same as training)

Edge Case Tests

• test_single_sample_batch: batch_size=1 works correctly
• test_zero_sequence_length: seq_len=0 handled gracefully
• test_large_batch_size: batch_size=64 works without memory errors

---

🔍 Debugging Tips

Test Fails with Shape Mismatch

cargo test -p ml test_forward_pass_shapes -- --nocapture

Look for: "Expected shape [batch, seq, d_model], got [batch, seq, 1]"

Test Fails with Dtype Error

cargo test -p ml test_all_tensors_dtype_f64 -- --nocapture

Look for: "Expected DType::F64, got DType::F32"

Test Fails with Broadcast Error

cargo test -p ml test_adam_optimizer_broadcasts -- --nocapture

Look for: "Incompatible dtypes for broadcast: F32 vs F64"

Training Loop Crashes

cargo test -p ml test_single_training_step -- --nocapture

Look for: "NaN detected in loss" or "Shape mismatch in loss computation"

---

📈 Performance Benchmarks

Operation	Target	Actual	Status
Test Suite Run	<10s	5-10s	✅ PASS
Single Test	<1s	<1s	✅ PASS
Forward Pass	<5μs	TBD	⏳ Pending Bug #18 fix
Training Step	<10ms	TBD	⏳ Pending Bug #18 fix

---

🛠️ Adding New Tests

Template for Shape Test

#[tokio::test]
async fn test_my_new_shape_validation() -> Result<()> {
    println!("🧪 Test: My New Shape Validation");

    let device = Device::Cpu;
    let config = minimal_test_config();
    let mut model = Mamba2SSM::new(config.clone(), &device)?;

    // Create input
    let input = Tensor::randn(0f64, 1.0, (batch_size, seq_len, d_model), &device)?;

    // Run operation
    let output = model.forward(&input)?;

    // Assert shape
    assert_eq!(output.dims(), &[expected_batch, expected_seq, expected_features],
        "Shape must be [batch, seq, features]");

    println!("✅ Test PASSED");
    Ok(())
}

Template for Dtype Test

#[tokio::test]
async fn test_my_dtype_validation() -> Result<()> {
    println!("🧪 Test: My Dtype Validation");

    let device = Device::Cpu;
    let tensor = Tensor::randn(0f64, 1.0, (2, 4), &device)?;

    // Assert dtype
    assert_eq!(tensor.dtype(), DType::F64,
        "Tensor must be F64, got {:?}", tensor.dtype());

    println!("✅ Test PASSED");
    Ok(())
}

---

🎓 Best Practices

1. Use Minimal Configs for Fast Tests

fn minimal_test_config() -> Mamba2Config {
    Mamba2Config {
        d_model: 16,        // Small for fast tests
        d_state: 4,         // Small state space
        expand: 2,          // Minimal expansion
        num_layers: 1,      // Single layer only
        batch_size: 2,      // Tiny batch
        seq_len: 8,         // Short sequences
        ...
    }
}

2. Assert Exact Shapes (Not Just Dimensions)

// ❌ BAD: Only checks number of dimensions
assert_eq!(output.dims().len(), 3);

// ✅ GOOD: Checks exact shape
assert_eq!(output.dims(), &[batch_size, seq_len, d_model],
    "Output must be [batch={}, seq={}, d_model={}]", batch_size, seq_len, d_model);

3. Test Edge Cases

// Test batch_size=1
// Test seq_len=0
// Test very large batch_size=64

4. Use Descriptive Test Names

// ❌ BAD: test_forward()
// ✅ GOOD: test_forward_pass_shapes()

---

📞 Support

Questions? See full documentation in AGENT_220_TDD_SHAPE_TESTS.md

Bug Reports? Run tests with --nocapture flag to see detailed output

Need Help? Check existing tests in /home/jgrusewski/Work/foxhunt/ml/tests/mamba2_shape_tests.rs

---

Last Updated: 2025-10-15 (Agent 220) Status: ✅ Tests created, 🔴 Bug #18 blocking 11/18 tests Next Step: Fix Bug #18 (LayerNorm F64 conversion)