monster/corridorkey-mrp-mlx
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix: reverse decoder concat order to match Torch (c4,c3,c2,c1)

Browse Source 
Torch decoder concatenates feature projections as [c4,c3,c2,c1] but
MLX was using [c1,c2,c3,c4]. The linear_fuse conv was trained with
the Torch order, producing scrambled features. Verified fix gives
correlation=1.0 against Torch on identical inputs.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
cmoyates 2026-03-01 09:44:41 -03:30
parent 81967ccc7c
commit 9cd800d37b
No known key found for this signature in database
GPG Key ID: F65E08480FDC22D2
1 changed files with 2 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
src/corridorkey_mlx/model/decoder.py
View File

@ -86,8 +86,8 @@ class DecoderHead(nn.Module):
x = up(x)
projected.append(x)
# Concatenate along channel dim (last dim in NHWC)
fused = mx.concatenate(projected, axis=-1) # (B, H/4, W/4, embed_dim*4)
# Concatenate in c4, c3, c2, c1 order to match Torch decoder
fused = mx.concatenate(projected[::-1], axis=-1) # (B, H/4, W/4, embed_dim*4)
fused = self.linear_fuse(fused)
fused = self.bn(fused)
fused = nn.relu(fused)