# BokehFlow Code Audit — Issues Found and Fixed

## CRITICAL BUGS (Will cause training failure or incorrect results)

### 1. ❌ Compositing double-multiplication bug (render_bokeh, line ~825)
**Problem:** `output = output + blurred * visible / (blurred_mask + 1e-6) * visible`  
This multiplies by `visible` TWICE — `blurred * visible / blurred_mask * visible` = wrong alpha compositing.  
**Fix:** `output = output + blurred / (blurred_mask + 1e-6) * visible`

### 2. ❌ CoC map computation doesn't handle focus distance == depth correctly  
**Problem:** When D == S₁ (pixel at focus distance), CoC should be exactly 0. The formula computes `abs(D - S₁)` which is correct, but the `S1.clamp(min=f+1.0)` can produce NaN gradients when `f` is a learnable parameter.  
**Fix:** Detach f from the clamp or use a fixed minimum.

### 3. ❌ BatchNorm in ConvStem will break at batch_size=1 during inference  
**Problem:** `nn.BatchNorm2d` computes running stats during training but fails with batch_size=1 if model is in training mode.  
**Fix:** Use `nn.GroupNorm(num_groups=8, num_channels=...)` or `nn.InstanceNorm2d` instead.

## STABILITY ISSUES (May cause NaN/Inf during training)

### 4. ⚠️ No gradient clipping mentioned in training config  
**Problem:** The GatedDeltaNet recurrence compounds matrix operations. Without gradient clipping, gradients can explode.  
**Fix:** Add `max_grad_norm=1.0` to training config.

### 5. ⚠️ Key L2-normalization — correct but needs epsilon  
**Problem:** `F.normalize(k, p=2, dim=-1)` can produce NaN if k is all zeros.  
**Fix:** Add eps: `k = F.normalize(k, p=2, dim=-1, eps=1e-8)`

### 6. ⚠️ State explosion risk  
**Problem:** The state update `state = a_t * (state - b_t * (state @ kk_t)) + b_t * vk_t` has matrix products that can grow unbounded if α≈1 and β≈0 for many steps.  
**Fix:** Add periodic state normalization: `state = state / (state.norm() + 1e-6).clamp(min=0.1)` every 256 steps.

### 7. ⚠️ Softplus depth output has no upper bound  
**Problem:** `nn.Softplus()` can output arbitrarily large values, causing CoC explosion.  
**Fix:** `depth = F.softplus(raw_depth).clamp(max=100.0)` (100 meters max).

## LOGICAL ISSUES

### 8. ⚠️ embed_dim mismatch for base variant  
**Problem:** `num_heads=6, head_dim=32` means inner_dim=192 but `embed_dim=192`, so the linear projections `to_qkv` project 192→3*192=576. This is correct but the output gate also projects 192→192. No bug but very heavy for base variant.

### 9. ⚠️ Direction fusion uses outputs before normalization  
**Problem:** The adaptive direction fusion `softmax(W_γ · [o_→;...])` operates on raw scan outputs, then the result is LayerNorm'd. The softmax inputs can have different scales per direction, potentially making one direction always dominate.  
**Fix:** Apply LayerNorm to each scan output BEFORE fusion, or use a temperature in the softmax.

### 10. ⚠️ TSP state shape mismatch  
**Problem:** `self.S_init` has shape `(1, num_heads, head_dim, head_dim)` but BiGDR returns a list of states (one per scan direction), not a single state. The propagate function iterates over block_states which are lists, not tensors.  
**Fix:** S_init should match the per-direction state shape, and propagation should handle the list structure properly.

## DATASET CONFIRMED COMPATIBLE ✅
RealBokeh has paired data:
- Input: `{split}/in/{id}_f22.JPG` (sharp, f/22)
- GT: `{split}/gt/{id}/{id}_f{fstop}.JPG` (variable bokeh)
- Metadata: `{split}/metadata/{id}.json` with focal_length, focus_plane_distance, target_avs

This maps perfectly to BokehFlow's inputs: image, f_number, focal_length_mm, focus_distance_m.