feat(phase4d): shape + parity tests, fix attention transpose

Fix MaskUnitAttention output transpose (0,2,3,1,4 → 0,3,2,1,4) to
match PyTorch transpose(1,3) token ordering for windowed attention.

Parity results (4 stages):
  Stage 0: max_abs 2.9e-4, mean 1.6e-5
  Stage 1: max_abs 1.3e-4, mean 1.3e-5
  Stage 2: max_abs 1.1e-2, mean 5.2e-5 (16 blocks, expected drift)
  Stage 3: max_abs 5.8e-4, mean 2.6e-5

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

src/corridorkey_mlx/model/hiera.py

@@ -289,9 +289,9 @@ class MaskUnitAttention(nn.Module):
         attn = mx.softmax(attn, axis=-1)
         x = attn @ v
 
-        # [B, heads, num_windows, tokens, head_dim] -> [B, num_windows, tokens, heads, head_dim]
-        # -> [B, N', dim_out]
-        x = mx.transpose(x, axes=(0, 2, 3, 1, 4))
+        # [B, heads, num_windows, tokens, head_dim] -> [B, tokens, num_windows, heads, head_dim]
+        # -> [B, N', dim_out]  (matches PyTorch transpose(1, 3))
+        x = mx.transpose(x, axes=(0, 3, 2, 1, 4))
         x = x.reshape(batch_size, -1, self.dim_out)
 
         return self.proj(x)

tests/test_hiera_stage_parity.py

@@ -0,0 +1,84 @@
+"""Parity tests: MLX Hiera backbone vs PyTorch reference (Phase 4).
+
+Loads golden input + encoder features from fixtures, runs MLX backbone
+with checkpoint weights, compares each stage output.
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+import mlx.core as mx
+import numpy as np
+import pytest
+
+from corridorkey_mlx.model.hiera import HieraBackbone
+from corridorkey_mlx.utils.layout import nchw_to_nhwc_np, nhwc_to_nchw_np
+
+FIXTURE_PATH = Path("reference/fixtures/golden.npz")
+CHECKPOINT_PATH = Path("checkpoints/corridorkey_mlx.safetensors")
+IMG_SIZE = 512
+NUM_STAGES = 4
+
+# Stage 2 runs 16 consecutive blocks — float32 drift accumulates on Metal vs CPU.
+# Mean error stays < 1e-4; max outliers can reach ~0.01 in the deepest stage.
+MAX_ABS_TOL = 2e-2
+
+
+def _skip_if_missing() -> None:
+    if not FIXTURE_PATH.exists():
+        pytest.skip("Fixture files not found — run dump_pytorch_reference.py first")
+    if not CHECKPOINT_PATH.exists():
+        pytest.skip("Checkpoint not found — run scripts/convert_weights.py first")
+
+
+@pytest.fixture(scope="module")
+def backbone_and_fixtures() -> (
+    tuple[list[mx.array], dict[str, np.ndarray]]
+):
+    """Load backbone once, return (mlx_features, fixtures)."""
+    _skip_if_missing()
+
+    fixtures = dict(np.load(FIXTURE_PATH))
+
+    # Load input: NCHW -> NHWC
+    input_nchw = fixtures["input"]
+    input_nhwc = mx.array(nchw_to_nhwc_np(input_nchw))
+
+    backbone = HieraBackbone(img_size=IMG_SIZE)
+    backbone.load_checkpoint(CHECKPOINT_PATH)
+
+    features = backbone(input_nhwc)
+    # materialize all features — mx.eval is MLX lazy evaluation, not Python eval
+    mx.eval(features)  # noqa: S307
+
+    return features, fixtures
+
+
+@pytest.mark.parametrize("stage_idx", range(NUM_STAGES))
+def test_stage_parity(
+    stage_idx: int,
+    backbone_and_fixtures: tuple[list[mx.array], dict[str, np.ndarray]],
+) -> None:
+    """MLX backbone stage output matches PyTorch within tolerance."""
+    features, fixtures = backbone_and_fixtures
+
+    expected_nchw = fixtures[f"encoder_feature_{stage_idx}"]
+    result_nhwc = features[stage_idx]
+    result_nchw = nhwc_to_nchw_np(np.array(result_nhwc))
+
+    assert result_nchw.shape == expected_nchw.shape, (
+        f"Stage {stage_idx} shape mismatch: {result_nchw.shape} vs {expected_nchw.shape}"
+    )
+
+    abs_err = np.abs(result_nchw - expected_nchw)
+    max_abs_err = float(np.max(abs_err))
+    mean_abs_err = float(np.mean(abs_err))
+    print(
+        f"\nStage {stage_idx} parity — "
+        f"max_abs: {max_abs_err:.6e}, mean_abs: {mean_abs_err:.6e}"
+    )
+
+    assert max_abs_err < MAX_ABS_TOL, (
+        f"Stage {stage_idx} max abs error {max_abs_err:.6e} exceeds tolerance {MAX_ABS_TOL}"
+    )

tests/test_hiera_stage_shapes.py

@@ -0,0 +1,46 @@
+"""Shape contract tests for Hiera backbone (Phase 4).
+
+No checkpoint needed — verifies structural correctness with random weights.
+"""
+
+from __future__ import annotations
+
+import mlx.core as mx
+import pytest
+
+from corridorkey_mlx.model.hiera import HieraBackbone, HieraPatchEmbed
+
+IMG_SIZE = 512
+NUM_FEATURES = 4
+EXPECTED_SHAPES = [
+    (1, 128, 128, 112),  # stride 4
+    (1, 64, 64, 224),  # stride 8
+    (1, 32, 32, 448),  # stride 16
+    (1, 16, 16, 896),  # stride 32
+]
+
+
+def test_patch_embed_output_shape() -> None:
+    """PatchEmbed produces (B, N, C) with N = (H/4)*(W/4)."""
+    patch_embed = HieraPatchEmbed()
+    x = mx.zeros((1, IMG_SIZE, IMG_SIZE, 4))
+    out = patch_embed(x)
+    expected_n = (IMG_SIZE // 4) * (IMG_SIZE // 4)
+    assert out.shape == (1, expected_n, 112)
+
+
+def test_backbone_returns_four_features() -> None:
+    """Backbone returns exactly 4 feature maps."""
+    backbone = HieraBackbone(img_size=IMG_SIZE)
+    x = mx.zeros((1, IMG_SIZE, IMG_SIZE, 4))
+    features = backbone(x)
+    assert len(features) == NUM_FEATURES
+
+
+@pytest.mark.parametrize("stage_idx", range(NUM_FEATURES))
+def test_feature_map_shape(stage_idx: int) -> None:
+    """Each stage feature map has correct (B, H, W, C) shape."""
+    backbone = HieraBackbone(img_size=IMG_SIZE)
+    x = mx.zeros((1, IMG_SIZE, IMG_SIZE, 4))
+    features = backbone(x)
+    assert features[stage_idx].shape == EXPECTED_SHAPES[stage_idx]