foxhunt/AGENT_172_QUICK_REFERENCE.md

AGENT 172: Quick Reference - MAMBA-2 Shape Bug

🎯 TL;DR

Problem: prepare_scan_input returns [8, 60, 1024] instead of [8, 60, 16]

Root Cause Hypothesis: B matrix has shape [1024, 1024] instead of [16, 1024]

Confidence: 85%

---

🚀 Quick Test

# 1. Clean rebuild
cargo clean -p ml
cargo build -p ml

# 2. Run test with debug output
cargo test -p ml test_mamba2_forward_pass --lib -- --nocapture 2>&1 | grep "AGENT 172 DEBUG"

# 3. Analyze output
# Look for: B shape, B.t() shape, Bu shape

---

📐 Expected Dimensions

Tensor	Expected Shape	Config
d_model	256	From config
d_state	16	From config
expand	4	From config
d_inner	1024	= d_model × expand
B	[16, 1024]	[d_state, d_inner]
B.t()	[1024, 16]	Transpose
input	[8, 60, 1024]	[batch, seq, d_inner]
Bu	[8, 60, 16]	input @ B.t()

---

🐛 If B Shape is Wrong

Scenario 1: B = [1024, 1024]

Fix Line 245:

// Check d_inner calculation at line 225
let d_inner = config.d_model * config.expand; // Must be 1024

// Verify B initialization at line 245
let B = Tensor::randn(0.0, 1.0, (config.d_state, d_inner), device)
// Should create [16, 1024], not [1024, 1024]

If d_inner is wrong, check:

1. config.d_model = 256 ✓
2. config.expand = 4 ✓
3. d_inner = 256 * 4 = 1024 ✓

Scenario 2: B = [16, 256]

Fix: Agent 168's fix not applied. Line 245 still uses config.d_model instead of d_inner:

// WRONG (old code)
let B = Tensor::randn(0.0, 1.0, (config.d_state, config.d_model), device)
// Creates [16, 256] ❌

// CORRECT (Agent 168 fix)
let B = Tensor::randn(0.0, 1.0, (config.d_state, d_inner), device)
// Creates [16, 1024] ✅

---

🔍 Debug Print Locations

All debug prints start with [AGENT 172 DEBUG] for easy grepping.

1. B Matrix Initialization (Line 251)

Shows shape immediately after creation.

2. forward_ssd_layer (Lines 617, 624, 630)

Shows input shape, B shape, B_discrete shape.

3. prepare_scan_input (Lines 701-716)

Shows all intermediate shapes during matmul.

---

🔧 Quick Fixes

Fix 1: Update Misleading Comment (Line 676)

// OLD
// FIXED: dt is [d_model] but B_cont is [d_state, d_model]

// NEW
// FIXED: dt is [d_model] but B_cont is [d_state, d_inner]

Fix 2: If MatMul is Broken

Try alternative matmul approaches in prepare_scan_input:

// Option 1: Explicit reshape
let B_t = B.t()?.reshape(&[1024, 16])?;
let Bu = input.matmul(&B_t)?;

// Option 2: broadcast_matmul
let Bu = input.broadcast_matmul(&B.t()?)?;

---

📊 Debug Output Template

Expected Output:

[AGENT 172 DEBUG] Layer 0 B matrix initialized: shape=[16, 1024], expected=[16, 1024]
[AGENT 172 DEBUG] Layer 1 B matrix initialized: shape=[16, 1024], expected=[16, 1024]
[AGENT 172 DEBUG] forward_ssd_layer layer 0: input shape=[8, 60, 1024]
[AGENT 172 DEBUG] forward_ssd_layer layer 0: B shape=[16, 1024]
[AGENT 172 DEBUG] forward_ssd_layer layer 0: B_discrete shape=[16, 1024]
[AGENT 172 DEBUG] prepare_scan_input shapes:
  input shape: [8, 60, 1024]
  B shape: [16, 1024]
  d_model: 256, d_inner: 1024, d_state: 16
  B.t() shape: [1024, 16]
  Bu shape: [8, 60, 16]
  Expected Bu shape: [batch=8, seq=60, d_state=16]

If Bug Persists (look for mismatched shapes):

[AGENT 172 DEBUG] Layer 0 B matrix initialized: shape=[1024, 1024], expected=[16, 1024]  ← WRONG!

---

🎯 Most Likely Scenarios (Ranked)

1. 70%: Agent 168's fix not compiled (old binary) - B still uses config.d_model instead of d_inner
2. 15%: d_inner calculation wrong - Something breaks d_model * expand
3. 10%: B gets corrupted during training - Gradient update reshapes B
4. 5%: Candle matmul bug - Returns wrong shape despite correct inputs

---

📝 Files Modified

• /home/jgrusewski/Work/foxhunt/ml/src/mamba/mod.rs
  • Line 251: B initialization debug
  • Line 617: forward_ssd_layer input debug
  • Line 624: forward_ssd_layer B debug
  • Line 630: forward_ssd_layer B_discrete debug
  • Lines 701-716: prepare_scan_input full debug trace

---

✅ Next Steps

1. Run test with debug output (command above)
2. Identify exact B shape from debug prints
3. Apply appropriate fix based on findings
4. Remove debug prints after fix validated
5. Document fix in AGENT_172_SUMMARY.md

---

Status: Debug infrastructure ready, waiting for test execution Created: 2025-10-15 Agent: 172