From cc2d31b435fb5fd78cccc94abc89cc678de9e69b Mon Sep 17 00:00:00 2001 From: cmoyates Date: Sun, 1 Mar 2026 05:51:36 -0330 Subject: [PATCH] feat(phase4d): shape + parity tests, fix attention transpose MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 --- src/corridorkey_mlx/model/hiera.py | 6 +-- tests/test_hiera_stage_parity.py | 84 ++++++++++++++++++++++++++++++ tests/test_hiera_stage_shapes.py | 46 ++++++++++++++++ 3 files changed, 133 insertions(+), 3 deletions(-) create mode 100644 tests/test_hiera_stage_parity.py create mode 100644 tests/test_hiera_stage_shapes.py diff --git a/src/corridorkey_mlx/model/hiera.py b/src/corridorkey_mlx/model/hiera.py index 047a5a7..9c7b591 100644 --- a/src/corridorkey_mlx/model/hiera.py +++ b/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) diff --git a/tests/test_hiera_stage_parity.py b/tests/test_hiera_stage_parity.py new file mode 100644 index 0000000..c9182b7 --- /dev/null +++ b/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}" + ) diff --git a/tests/test_hiera_stage_shapes.py b/tests/test_hiera_stage_shapes.py new file mode 100644 index 0000000..0780d39 --- /dev/null +++ b/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]