# AGENT 182 QUICK FIX: parallel_prefix_scan Shape Bug

**Mission**: Fix `parallel_prefix_scan` to preserve `[batch, seq, d_state]` shape

**Priority**: 🔴 **CRITICAL** - Blocking all MAMBA-2 training (0/7 tests passing)

---

## 🎯 The Bug

**File**: `/home/jgrusewski/Work/foxhunt/ml/src/mamba/scan_algorithms.rs:111`

**Function**: `parallel_prefix_scan`

**Problem**: Returns `[batch, seq, d_inner]` instead of `[batch, seq, d_state]`

**Impact**: Causes shape mismatch at line 633 of `/home/jgrusewski/Work/foxhunt/ml/src/mamba/mod.rs`:

```rust
let output = scanned_states.matmul(&C.t()?)?;
// ERROR: [8, 60, 1024] @ [1024, 16] - dimension mismatch!
```

---

## 🔍 Root Cause

### Expected Behavior

```
Input to parallel_prefix_scan:  [8, 60, 16] (d_state)
Output from parallel_prefix_scan: [8, 60, 16] (preserve shape)
```

### Actual Behavior

```
Input to parallel_prefix_scan:  [8, 60, 16] (d_state)
Output from parallel_prefix_scan: [8, 60, 1024] (d_inner) ❌ WRONG!
```

### Where the Bug Occurs

The scan algorithm is likely using the **wrong tensor** in one of these functions:

1. `sequential_scan` (line 148)
2. `block_parallel_scan` (called from line 124)

**Hypothesis**: One of these functions is using the **original input** (`[*, *, d_inner]`) instead of the **scan input** (`[*, *, d_state]`).

---

## 🔧 Investigation Steps

### Step 1: Check sequential_scan

```bash
# Search for where the result tensor is created in sequential_scan
grep -A 30 "fn sequential_scan" ml/src/mamba/scan_algorithms.rs
```

**Look for**:
- Result tensor creation
- Shape used for result allocation
- Which tensor is being scanned (should be `input` parameter, not anything else)

### Step 2: Check block_parallel_scan

```bash
# Search for block_parallel_scan implementation
grep -A 50 "fn block_parallel_scan" ml/src/mamba/scan_algorithms.rs
```

**Look for**:
- Block size calculations using wrong dimensions
- Result tensor shape allocation
- Concatenation operations that might expand dimensions

### Step 3: Look for d_inner references

```bash
# Check if scan_algorithms.rs incorrectly references d_inner
grep -n "d_inner\|1024" ml/src/mamba/scan_algorithms.rs
```

**Expected**: NO references to `d_inner` or hardcoded `1024` in scan_algorithms.rs

---

## 🎯 Likely Fix

### Scenario A: Using Wrong Tensor

If the scan is using `self.state.hidden` or `input_projection` output instead of the `input` parameter:

```rust
// WRONG:
let result = self.scan(self.hidden_state)?; // Uses d_inner dimension

// CORRECT:
let result = self.scan(input)?; // Uses d_state dimension from parameter
```

### Scenario B: Wrong Result Shape Allocation

If the result tensor is allocated with wrong dimensions:

```rust
// WRONG:
let result = Tensor::zeros((batch_size, seq_len, d_inner), ...)?;

// CORRECT:
let result = Tensor::zeros((batch_size, seq_len, input.dim(2)?), ...)?;
```

### Scenario C: Accumulator Shape Bug

If the accumulator in `sequential_scan` is using wrong shape:

```rust
// WRONG:
let mut accumulator = Tensor::zeros((batch_size, 1, d_inner), ...)?;

// CORRECT:
let mut accumulator = input.narrow(0, 0, 1)?.narrow(1, 0, 1)?; // Use input shape
```

---

## 📝 Files to Modify

**Primary**:
- `/home/jgrusewski/Work/foxhunt/ml/src/mamba/scan_algorithms.rs`

**Verify**:
- `/home/jgrusewski/Work/foxhunt/ml/src/mamba/mod.rs` (no changes needed, already correct)

---

## ✅ Success Criteria

After fix, run:

```bash
cargo test -p ml --test e2e_mamba2_training --features cuda
```

**Expected**:
```
test result: ok. 7 passed; 0 failed
```

**Test that will pass first**: `test_mamba2_simple_forward_pass`

**Shape trace should show**:
```
scan_input:      [8, 60, 16] ✅
scanned_states:  [8, 60, 16] ✅ (not [8, 60, 1024])
output:          [8, 60, 1024] ✅
```

---

## 🚨 Critical Notes

1. **DO NOT modify B/C matrix shapes** - They are already correct!
2. **DO NOT modify prepare_scan_input** - It's working correctly!
3. **ONLY fix the scan algorithm** - Shape should be preserved

---

## 📊 Test Configuration

```rust
d_model: 256
d_state: 16
expand: 4
d_inner: 1024 (256 * 4)

B: [16, 1024] (d_state × d_inner) ✅
C: [1024, 16] (d_inner × d_state) ✅
scan_input: [8, 60, 16] ✅
scanned_states: [8, 60, 16] ← FIX THIS (currently [8, 60, 1024])
```

---

## 🔬 Debugging Commands

```bash
# Run single test with full output
cargo test -p ml test_mamba2_simple_forward_pass --features cuda -- --nocapture

# Check scan_algorithms.rs for dimension bugs
rg "d_inner|1024" ml/src/mamba/scan_algorithms.rs

# Look for tensor shape allocations
rg "Tensor::zeros|Tensor::ones" ml/src/mamba/scan_algorithms.rs
```

---

## ⏱️ Estimated Fix Time

**30-60 minutes** (scan algorithm is isolated module)

**Confidence**: ✅ High - Root cause clearly identified, fix is localized

---

**End of Quick Fix Guide**