From 838bd0dd390efd039fafbce839cbf3addad2861e Mon Sep 17 00:00:00 2001 From: cmoyates Date: Tue, 3 Mar 2026 09:58:53 -0330 Subject: [PATCH] =?UTF-8?q?refactor(tests):=20restructure=20suite=20?= =?UTF-8?q?=E2=80=94=2015=20files=20to=208,=20consolidate=20parity?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Add conftest.py: shared paths, tolerances, skip markers, fixtures - Consolidate imports/shapes/smoke/forward → test_model_contract.py - Consolidate all parity (decoder/refiner/backbone/e2e) → test_parity.py - Fold 2048 slow test into test_engine.py - Rework test_conversion.py to use public convert_checkpoint() API - Simplify test_weights.py: drop CLI/env var tests, keep checksum+config - Rename tiling/compilation test files for consistency - Delete 10 redundant files 80 tests collected (75 pass, 4 skip, 1 deselected) Co-Authored-By: Claude Opus 4.6 --- tests/conftest.py | 71 ++++++ ...ger_consistency.py => test_compilation.py} | 0 tests/test_conversion.py | 222 +++++------------- tests/test_decoder_parity.py | 150 ------------ tests/test_end_to_end_parity.py | 94 -------- tests/test_end_to_end_smoke.py | 66 ------ tests/test_engine.py | 48 +++- tests/test_greenformer_forward.py | 77 ------ tests/test_hiera_stage_parity.py | 79 ------- tests/test_hiera_stage_shapes.py | 46 ---- tests/test_import.py | 12 - tests/test_model_contract.py | 194 +++++++++++++++ tests/test_parity.py | 140 +++++++++++ tests/test_reference_fixtures.py | 130 ---------- tests/test_refiner_parity.py | 154 ------------ tests/test_smoke_2048.py | 78 ------ ...t_tiling_consistency.py => test_tiling.py} | 0 tests/test_weights.py | 64 +---- 18 files changed, 509 insertions(+), 1116 deletions(-) create mode 100644 tests/conftest.py rename tests/{test_compiled_vs_eager_consistency.py => test_compilation.py} (100%) delete mode 100644 tests/test_decoder_parity.py delete mode 100644 tests/test_end_to_end_parity.py delete mode 100644 tests/test_end_to_end_smoke.py delete mode 100644 tests/test_greenformer_forward.py delete mode 100644 tests/test_hiera_stage_parity.py delete mode 100644 tests/test_hiera_stage_shapes.py delete mode 100644 tests/test_import.py create mode 100644 tests/test_model_contract.py create mode 100644 tests/test_parity.py delete mode 100644 tests/test_reference_fixtures.py delete mode 100644 tests/test_refiner_parity.py delete mode 100644 tests/test_smoke_2048.py rename tests/{test_tiling_consistency.py => test_tiling.py} (100%) diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 0000000..5e49ca6 --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,71 @@ +"""Shared test fixtures, paths, tolerances, and skip markers.""" + +from __future__ import annotations + +from pathlib import Path + +import mlx.core as mx +import numpy as np +import pytest + +from corridorkey_mlx.model.corridorkey import GreenFormer + +# --------------------------------------------------------------------------- +# Paths +# --------------------------------------------------------------------------- +GOLDEN_PATH = Path("reference/fixtures/golden.npz") +GOLDEN_WEIGHTS_PATH = Path("reference/fixtures/golden_weights.npz") +MLX_CHECKPOINT_PATH = Path("checkpoints/corridorkey_mlx.safetensors") +PT_CHECKPOINT_PATH = Path("checkpoints/CorridorKey_v1.0.pth") + +# --------------------------------------------------------------------------- +# Tolerances +# --------------------------------------------------------------------------- +PARITY_TOL_TIGHT = 1e-4 +PARITY_TOL_BACKBONE = 2e-2 +PARITY_TOL_E2E = 1e-3 + +# --------------------------------------------------------------------------- +# Skip markers +# --------------------------------------------------------------------------- +has_golden = pytest.mark.skipif(not GOLDEN_PATH.exists(), reason="golden.npz not found") +has_checkpoint = pytest.mark.skipif( + not MLX_CHECKPOINT_PATH.exists(), reason="MLX checkpoint not found" +) +has_pt_checkpoint = pytest.mark.skipif( + not PT_CHECKPOINT_PATH.exists(), reason="PyTorch checkpoint not found" +) + +# --------------------------------------------------------------------------- +# Shared fixtures +# --------------------------------------------------------------------------- +IMG_SIZE = 512 +SMALL_IMG_SIZE = 256 + + +@pytest.fixture(scope="module") +def golden_fixtures() -> dict[str, np.ndarray]: + """Load golden.npz fixture file (PyTorch NCHW format).""" + if not GOLDEN_PATH.exists(): + pytest.skip("golden.npz not found") + return dict(np.load(GOLDEN_PATH)) + + +@pytest.fixture(scope="module") +def random_model() -> GreenFormer: + """GreenFormer with random weights at 512.""" + model = GreenFormer(img_size=IMG_SIZE) + model.eval() + # mx.eval is MLX array materialization, not Python eval() + mx.eval(model.parameters()) # noqa: S307 + return model + + +@pytest.fixture(scope="module") +def loaded_model() -> GreenFormer: + """GreenFormer with checkpoint weights at 512.""" + if not MLX_CHECKPOINT_PATH.exists(): + pytest.skip("MLX checkpoint not found") + model = GreenFormer(img_size=IMG_SIZE) + model.load_checkpoint(MLX_CHECKPOINT_PATH) + return model diff --git a/tests/test_compiled_vs_eager_consistency.py b/tests/test_compilation.py similarity index 100% rename from tests/test_compiled_vs_eager_consistency.py rename to tests/test_compilation.py diff --git a/tests/test_conversion.py b/tests/test_conversion.py index 85b225a..1384d37 100644 --- a/tests/test_conversion.py +++ b/tests/test_conversion.py @@ -1,185 +1,83 @@ -"""Tests: weight conversion PyTorch → MLX (Phase 3). +"""Weight conversion tests through public API. -Validates key mapping, shape transforms, and round-trip integrity. +Reworked from internal convert_state_dict tests to use convert_checkpoint(). """ from __future__ import annotations -from pathlib import Path from typing import TYPE_CHECKING import numpy as np import pytest +from corridorkey_mlx.convert.converter import convert_checkpoint +from corridorkey_mlx.model.corridorkey import GreenFormer + +from .conftest import PT_CHECKPOINT_PATH, has_pt_checkpoint + if TYPE_CHECKING: - from collections import OrderedDict + from pathlib import Path -from corridorkey_mlx.convert.converter import ( - REFINER_STEM_MAP, - SKIP_SUFFIXES, - convert_state_dict, - load_pytorch_checkpoint, -) +EXPECTED_KEY_COUNT = 365 -CHECKPOINT_PATH = Path("checkpoints/CorridorKey_v1.0.pth") -SAFETENSORS_PATH = Path("checkpoints/corridorkey_mlx.safetensors") + +def _torch_available() -> bool: + try: + import torch # noqa: F401 + + return True + except ImportError: + return False + + +has_torch = pytest.mark.skipif(not _torch_available(), reason="torch not installed") @pytest.fixture(scope="module") -def checkpoint_data() -> tuple[dict[str, np.ndarray], OrderedDict[str, np.ndarray], list]: - """Load and convert checkpoint once for all tests in this module.""" - if not CHECKPOINT_PATH.exists(): - pytest.skip("Checkpoint not found — need CorridorKey_v1.0.pth") - state_dict = load_pytorch_checkpoint(CHECKPOINT_PATH) - converted, diagnostics = convert_state_dict(state_dict) - return state_dict, converted, diagnostics +def converted_path(tmp_path_factory: pytest.TempPathFactory) -> Path: + """Convert PT checkpoint to safetensors in temp dir.""" + if not PT_CHECKPOINT_PATH.exists(): + pytest.skip("PyTorch checkpoint not found") + if not _torch_available(): + pytest.skip("torch not installed") + output = tmp_path_factory.mktemp("weights") / "converted.safetensors" + convert_checkpoint(PT_CHECKPOINT_PATH, output) + return output -# --------------------------------------------------------------------------- -# Key mapping completeness -# --------------------------------------------------------------------------- -class TestKeyMapping: - """Every PyTorch key maps to an MLX key with no orphans.""" - - def test_no_orphan_source_keys(self, checkpoint_data) -> None: - """All source keys are either mapped or explicitly skipped.""" - state_dict, _, diagnostics = checkpoint_data - - mapped_src_keys = {r.source_key for r in diagnostics} - skipped_keys = {k for k in state_dict if any(k.endswith(s) for s in SKIP_SUFFIXES)} - accounted_for = mapped_src_keys | skipped_keys - - orphans = set(state_dict.keys()) - accounted_for - assert orphans == set(), f"Orphan source keys: {orphans}" - - def test_no_duplicate_dest_keys(self, checkpoint_data) -> None: - """No two source keys map to the same destination key.""" - _, _, diagnostics = checkpoint_data - - dest_keys = [r.dest_key for r in diagnostics] - assert len(dest_keys) == len(set(dest_keys)), "Duplicate destination keys found" - - def test_refiner_stem_remapped(self, checkpoint_data) -> None: - """Refiner stem Sequential keys correctly remapped to named attrs.""" - _, converted, _ = checkpoint_data - - for pt_key, mlx_key in REFINER_STEM_MAP.items(): - assert mlx_key in converted, f"Expected remapped key {mlx_key} not found" - assert pt_key not in converted, f"Original key {pt_key} should not be in output" - - def test_num_batches_tracked_dropped(self, checkpoint_data) -> None: - """BatchNorm num_batches_tracked keys are not in output.""" - _, converted, _ = checkpoint_data - - for key in converted: - assert "num_batches_tracked" not in key, f"Found num_batches_tracked: {key}" - - def test_expected_key_count(self, checkpoint_data) -> None: - """Output has expected number of keys (367 source - 2 skipped = 365).""" - state_dict, converted, _ = checkpoint_data - - source_count = len(state_dict) - skip_count = sum(1 for k in state_dict if any(k.endswith(s) for s in SKIP_SUFFIXES)) - expected = source_count - skip_count - assert len(converted) == expected, f"Expected {expected} keys, got {len(converted)}" - - -# --------------------------------------------------------------------------- -# Conv weight transpose -# --------------------------------------------------------------------------- -class TestConvWeightTranspose: - """Conv weights correctly transposed from NCHW → NHWC.""" - - def test_patch_embed_conv_shape(self, checkpoint_data) -> None: - """Patch embed conv transposed: (112,4,7,7) → (112,7,7,4).""" - _, converted, _ = checkpoint_data - - key = "encoder.model.patch_embed.proj.weight" - assert key in converted - assert converted[key].shape == (112, 7, 7, 4) - - def test_decoder_conv_shapes(self, checkpoint_data) -> None: - """Decoder conv weights transposed correctly.""" - _, converted, _ = checkpoint_data - - # linear_fuse: (256, 1024, 1, 1) → (256, 1, 1, 1024) - for prefix in ("alpha_decoder", "fg_decoder"): - fuse_key = f"{prefix}.linear_fuse.weight" - assert converted[fuse_key].shape == (256, 1, 1, 1024) - - def test_refiner_conv_shapes(self, checkpoint_data) -> None: - """Refiner conv weights transposed correctly.""" - _, converted, _ = checkpoint_data - - # stem_conv: (64, 7, 3, 3) → (64, 3, 3, 7) - assert converted["refiner.stem_conv.weight"].shape == (64, 3, 3, 7) - - # res block convs: (64, 64, 3, 3) → (64, 3, 3, 64) - for i in range(1, 5): - for j in range(1, 3): - key = f"refiner.res{i}.conv{j}.weight" - assert converted[key].shape == (64, 3, 3, 64), f"Wrong shape for {key}" - - # final: (4, 64, 1, 1) → (4, 1, 1, 64) - assert converted["refiner.final.weight"].shape == (4, 1, 1, 64) - - def test_linear_weights_unchanged(self, checkpoint_data) -> None: - """Linear (2D) weights are not transposed.""" - _, _, diagnostics = checkpoint_data - - passthrough_2d = [ - r for r in diagnostics if r.transform == "passthrough" and len(r.source_shape) == 2 - ] - for record in passthrough_2d: - assert record.source_shape == record.dest_shape, ( - f"{record.source_key}: shape changed despite passthrough" - ) - - -# --------------------------------------------------------------------------- -# 4-channel first conv -# --------------------------------------------------------------------------- -class TestFirstConv4Channel: - """Patched 4-channel first conv preserved during conversion.""" - - def test_input_channels_preserved(self, checkpoint_data) -> None: - """Patch embed conv has 4 input channels (RGB + alpha hint).""" - _, converted, _ = checkpoint_data - - weight = converted["encoder.model.patch_embed.proj.weight"] - # MLX layout: (O, H, W, I) — I is last dim - input_channels = weight.shape[-1] - assert input_channels == 4, f"Expected 4 input channels, got {input_channels}" - - def test_values_preserved(self, checkpoint_data) -> None: - """Conv values are preserved (only transposed, not modified).""" - state_dict, converted, _ = checkpoint_data - - pt_weight = state_dict["encoder.model.patch_embed.proj.weight"] - mlx_weight = converted["encoder.model.patch_embed.proj.weight"] - - # Transpose back to PyTorch layout and compare - roundtrip = np.transpose(mlx_weight, (0, 3, 1, 2)) - np.testing.assert_array_equal(roundtrip, pt_weight) - - -# --------------------------------------------------------------------------- -# Safetensors output validation -# --------------------------------------------------------------------------- -class TestSafetensorsOutput: - """Validate saved safetensors file matches conversion output.""" - - def test_safetensors_loadable(self, checkpoint_data) -> None: - """Saved safetensors file can be loaded and has correct keys.""" - if not SAFETENSORS_PATH.exists(): - pytest.skip("Run convert_weights.py first") +@has_pt_checkpoint +@has_torch +class TestConvertCheckpoint: + def test_output_file_exists(self, converted_path: Path) -> None: + assert converted_path.exists() + def test_safetensors_loadable(self, converted_path: Path) -> None: from safetensors.numpy import load_file - loaded = load_file(str(SAFETENSORS_PATH)) + weights = load_file(str(converted_path)) + assert len(weights) > 0 - _, converted, _ = checkpoint_data + def test_key_count(self, converted_path: Path) -> None: + from safetensors.numpy import load_file - assert set(loaded.keys()) == set(converted.keys()) - for key in converted: - np.testing.assert_array_equal(loaded[key], converted[key]) + weights = load_file(str(converted_path)) + assert len(weights) == EXPECTED_KEY_COUNT + + def test_model_loads_and_runs(self, converted_path: Path) -> None: + """Roundtrip: convert -> load into GreenFormer -> forward succeeds.""" + import mlx.core as mx + + model = GreenFormer(img_size=256) + model.load_checkpoint(converted_path) + + x = mx.random.normal((1, 256, 256, 4)) + out = model(x) + # mx.eval is MLX array materialization, not Python eval() + mx.eval(out) # noqa: S307 + + assert out["alpha_final"].shape == (1, 256, 256, 1) + assert out["fg_final"].shape == (1, 256, 256, 3) + + # No NaN + alpha = np.array(out["alpha_final"]) + assert not np.isnan(alpha).any() diff --git a/tests/test_decoder_parity.py b/tests/test_decoder_parity.py deleted file mode 100644 index 41aa5d0..0000000 --- a/tests/test_decoder_parity.py +++ /dev/null @@ -1,150 +0,0 @@ -"""Parity tests: MLX decoder heads vs PyTorch reference (Phase 2). - -Uses saved backbone features -> runs MLX decoder -> compares -against saved PyTorch coarse predictions. -""" - -from __future__ import annotations - -from pathlib import Path - -import mlx.core as mx -import numpy as np -import pytest - -from corridorkey_mlx.model.decoder import DecoderHead -from corridorkey_mlx.utils.layout import conv_weight_pt_to_mlx, nchw_to_nhwc_np, nhwc_to_nchw_np - -FIXTURE_PATH = Path("reference/fixtures/golden.npz") -WEIGHTS_PATH = Path("reference/fixtures/golden_weights.npz") - -BACKBONE_CHANNELS = [112, 224, 448, 896] -EMBED_DIM = 256 -MAX_ABS_TOL = 1e-4 # relaxed slightly for float32 Metal vs CPU differences - - -def _skip_if_missing() -> None: - if not FIXTURE_PATH.exists() or not WEIGHTS_PATH.exists(): - pytest.skip("Fixture files not found — run dump_pytorch_reference.py first") - - -def _load_decoder_weights( - prefix: str, - weights: dict[str, np.ndarray], -) -> dict[str, mx.array]: - """Convert PyTorch decoder state_dict to MLX parameter dict. - - Maps PyTorch keys (e.g. 'alpha_decoder.linear_c1.proj.weight') to - MLX nested keys (e.g. 'linear_c1.proj.weight'), transposing conv weights. - """ - params: dict[str, mx.array] = {} - for pt_key, value in weights.items(): - if not pt_key.startswith(prefix): - continue - mlx_key = pt_key[len(prefix) :] - - # Conv2d weights: (O,I,H,W) -> (O,H,W,I) - if "linear_fuse.weight" in mlx_key or "classifier.weight" in mlx_key: - value = conv_weight_pt_to_mlx(value) - - # BatchNorm: running_mean -> running_mean, running_var -> running_var - # PyTorch uses num_batches_tracked which MLX doesn't need - if "num_batches_tracked" in mlx_key: - continue - - params[mlx_key] = mx.array(value) - - return params - - -def _params_to_nested(flat: dict[str, mx.array]) -> dict: - """Convert flat dot-separated keys to nested dict for mlx load_weights.""" - nested: dict = {} - for key, value in flat.items(): - parts = key.split(".") - current = nested - for part in parts[:-1]: - if part not in current: - current[part] = {} - current = current[part] - current[parts[-1]] = value - return nested - - -def _build_and_load_decoder( - output_dim: int, - prefix: str, - weights: dict[str, np.ndarray], -) -> DecoderHead: - """Build an MLX DecoderHead and load converted weights.""" - decoder = DecoderHead(BACKBONE_CHANNELS, EMBED_DIM, output_dim) - flat_params = _load_decoder_weights(prefix, weights) - nested_params = _params_to_nested(flat_params) - decoder.load_weights(list(_flatten_nested(nested_params))) - decoder.eval() # use running stats for BatchNorm (not batch stats) - mx.eval(decoder.parameters()) # noqa: S307 — mx.eval is MLX lazy eval, not Python eval - return decoder - - -def _flatten_nested(d: dict, prefix: str = "") -> list[tuple[str, mx.array]]: - """Flatten nested dict to list of (dotted_key, array) pairs.""" - items: list[tuple[str, mx.array]] = [] - for k, v in d.items(): - full_key = f"{prefix}.{k}" if prefix else k - if isinstance(v, dict): - items.extend(_flatten_nested(v, full_key)) - else: - items.append((full_key, v)) - return items - - -def _run_decoder_parity( - output_dim: int, - prefix: str, - expected_key: str, -) -> None: - """Run a decoder parity test.""" - _skip_if_missing() - - fixtures = dict(np.load(FIXTURE_PATH)) - weights = dict(np.load(WEIGHTS_PATH)) - - # Load features in NHWC - features_nhwc = [mx.array(nchw_to_nhwc_np(fixtures[f"encoder_feature_{i}"])) for i in range(4)] - - decoder = _build_and_load_decoder(output_dim, prefix, weights) - result_nhwc = decoder(features_nhwc) - mx.eval(result_nhwc) # noqa: S307 — mx.eval is MLX lazy eval, not Python eval - - # Convert result back to NCHW for comparison - result_nchw = nhwc_to_nchw_np(np.array(result_nhwc)) - expected = fixtures[expected_key] - - max_abs_err = float(np.max(np.abs(result_nchw - expected))) - mean_abs_err = float(np.mean(np.abs(result_nchw - expected))) - print(f"\n{prefix} parity — max_abs: {max_abs_err:.6e}, mean_abs: {mean_abs_err:.6e}") - - assert result_nchw.shape == expected.shape, ( - f"Shape mismatch: {result_nchw.shape} vs {expected.shape}" - ) - assert max_abs_err < MAX_ABS_TOL, ( - f"Max abs error {max_abs_err:.6e} exceeds tolerance {MAX_ABS_TOL}" - ) - - -def test_alpha_decoder_parity() -> None: - """MLX alpha decoder matches PyTorch within tolerance.""" - _run_decoder_parity( - output_dim=1, - prefix="alpha_decoder.", - expected_key="alpha_logits", - ) - - -def test_fg_decoder_parity() -> None: - """MLX foreground decoder matches PyTorch within tolerance.""" - _run_decoder_parity( - output_dim=3, - prefix="fg_decoder.", - expected_key="fg_logits", - ) diff --git a/tests/test_end_to_end_parity.py b/tests/test_end_to_end_parity.py deleted file mode 100644 index 666a0ce..0000000 --- a/tests/test_end_to_end_parity.py +++ /dev/null @@ -1,94 +0,0 @@ -"""End-to-end parity tests: MLX GreenFormer vs PyTorch golden reference (Phase 5). - -Loads golden.npz (PyTorch NCHW), runs same input through MLX model, -compares all intermediate and final outputs. -""" - -from __future__ import annotations - -from pathlib import Path - -import mlx.core as mx -import numpy as np -import pytest - -from corridorkey_mlx.model.corridorkey import GreenFormer -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 - -# Tolerance tiers — coarse path inherits backbone drift, refiner adds more. -# Backbone stages have up to ~0.01 max abs err (16 sequential blocks). -# Final outputs pass through sigmoid which compresses errors. -TOLERANCES: dict[str, float] = { - "alpha_logits": 1e-3, - "fg_logits": 1e-3, - "alpha_logits_up": 1e-3, - "fg_logits_up": 1e-3, - "alpha_coarse": 1e-4, - "fg_coarse": 1e-4, - "delta_logits": 2e-3, - "alpha_final": 2e-4, - "fg_final": 2e-4, -} - - -def _skip_if_missing() -> None: - if not FIXTURE_PATH.exists(): - pytest.skip("golden.npz 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 model_outputs_and_fixtures() -> tuple[dict[str, mx.array], dict[str, np.ndarray]]: - """Load model, run forward pass with golden input, return (mlx_outputs, fixtures).""" - _skip_if_missing() - - fixtures = dict(np.load(FIXTURE_PATH)) - - # Load input: NCHW -> NHWC - input_nhwc = mx.array(nchw_to_nhwc_np(fixtures["input"])) - - model = GreenFormer(img_size=IMG_SIZE) - model.load_checkpoint(CHECKPOINT_PATH) - - outputs = model(input_nhwc) - # materialize — mx.eval is MLX lazy graph evaluation, not Python eval - mx.eval(outputs) # noqa: S307 - - return outputs, fixtures - - -@pytest.mark.parametrize("key", list(TOLERANCES.keys())) -def test_parity( - key: str, - model_outputs_and_fixtures: tuple[dict[str, mx.array], dict[str, np.ndarray]], -) -> None: - """MLX output matches PyTorch golden reference within tolerance.""" - outputs, fixtures = model_outputs_and_fixtures - - assert key in outputs, f"Missing output key: {key}" - assert key in fixtures, f"Missing fixture key: {key}" - - # MLX output is NHWC, fixture is NCHW — convert MLX to NCHW for comparison - mlx_nchw = nhwc_to_nchw_np(np.array(outputs[key])) - expected_nchw = fixtures[key] - - assert mlx_nchw.shape == expected_nchw.shape, ( - f"{key} shape mismatch: MLX {mlx_nchw.shape} vs PyTorch {expected_nchw.shape}" - ) - - abs_err = np.abs(mlx_nchw - expected_nchw) - max_abs = float(np.max(abs_err)) - mean_abs = float(np.mean(abs_err)) - tol = TOLERANCES[key] - - print( - f"\n{key:<20s} | shape={mlx_nchw.shape} | " - f"max_abs={max_abs:.6e} | mean_abs={mean_abs:.6e} | tol={tol:.1e}" - ) - - assert max_abs < tol, f"{key}: max abs error {max_abs:.6e} exceeds tolerance {tol:.1e}" diff --git a/tests/test_end_to_end_smoke.py b/tests/test_end_to_end_smoke.py deleted file mode 100644 index 179e9e6..0000000 --- a/tests/test_end_to_end_smoke.py +++ /dev/null @@ -1,66 +0,0 @@ -"""End-to-end smoke test for GreenFormer (Phase 5). - -Verifies the full pipeline works with random weights — no checkpoint needed. -Tests model construction, forward pass, and basic output sanity. -""" - -from __future__ import annotations - -import mlx.core as mx -import numpy as np -import pytest - -from corridorkey_mlx.io.image import ( - postprocess_alpha, - postprocess_foreground, - preprocess, -) -from corridorkey_mlx.model.corridorkey import GreenFormer - -IMG_SIZE = 256 # smaller for speed - - -@pytest.fixture(scope="module") -def model() -> GreenFormer: - return GreenFormer(img_size=IMG_SIZE) - - -def test_forward_with_preprocessed_input(model: GreenFormer) -> None: - """Full pipeline: numpy image -> preprocess -> model -> postprocess.""" - rgb = np.random.rand(IMG_SIZE, IMG_SIZE, 3).astype(np.float32) - alpha_hint = np.random.rand(IMG_SIZE, IMG_SIZE, 1).astype(np.float32) - - x = preprocess(rgb, alpha_hint) - assert x.shape == (1, IMG_SIZE, IMG_SIZE, 4) - - out = model(x) - # materialize — mx.eval is MLX's lazy graph evaluation, not Python eval - mx.eval(out) - - alpha = postprocess_alpha(out["alpha_final"]) - fg = postprocess_foreground(out["fg_final"]) - - assert alpha.shape == (IMG_SIZE, IMG_SIZE) - assert alpha.dtype == np.uint8 - assert fg.shape == (IMG_SIZE, IMG_SIZE, 3) - assert fg.dtype == np.uint8 - - -def test_deterministic_output(model: GreenFormer) -> None: - """Same input produces same output.""" - mx.random.seed(123) - x = mx.random.normal((1, IMG_SIZE, IMG_SIZE, 4)) - # materialize — mx.eval is MLX's lazy graph evaluation, not Python eval - mx.eval(x) - - out1 = model(x) - mx.eval(out1) - out2 = model(x) - mx.eval(out2) - - for key in out1: - np.testing.assert_array_equal( - np.array(out1[key]), - np.array(out2[key]), - err_msg=f"{key} not deterministic", - ) diff --git a/tests/test_engine.py b/tests/test_engine.py index 287945c..373389c 100644 --- a/tests/test_engine.py +++ b/tests/test_engine.py @@ -2,8 +2,6 @@ from __future__ import annotations -from pathlib import Path - import numpy as np import pytest @@ -13,9 +11,7 @@ from corridorkey_mlx.engine import ( _validate_mask, ) -CHECKPOINT = Path("checkpoints/corridorkey_mlx.safetensors") -HAS_CHECKPOINT = CHECKPOINT.exists() - +from .conftest import MLX_CHECKPOINT_PATH, has_checkpoint # --------------------------------------------------------------------------- # Input validation (no checkpoint needed) @@ -87,14 +83,14 @@ class TestEngineInit: # --------------------------------------------------------------------------- -@pytest.mark.skipif(not HAS_CHECKPOINT, reason="Checkpoint not available") +@has_checkpoint class TestEngineIntegration: """Integration tests that load the real model.""" @pytest.fixture(scope="class") def engine(self) -> CorridorKeyMLXEngine: return CorridorKeyMLXEngine( - checkpoint_path=CHECKPOINT, + checkpoint_path=MLX_CHECKPOINT_PATH, img_size=512, compile=False, ) @@ -128,10 +124,42 @@ class TestEngineIntegration: result = engine.process_frame(image, mask) assert "alpha" in result - def test_refiner_scale_zero_returns_coarse( - self, engine: CorridorKeyMLXEngine - ) -> None: + def test_refiner_scale_zero_returns_coarse(self, engine: CorridorKeyMLXEngine) -> None: image = np.random.default_rng(42).integers(0, 256, (64, 64, 3), dtype=np.uint8) mask = np.random.default_rng(42).integers(0, 256, (64, 64), dtype=np.uint8) result = engine.process_frame(image, mask, refiner_scale=0.0) assert result["alpha"].shape == (64, 64) + + +# --------------------------------------------------------------------------- +# Full 2048 inference (slow, requires checkpoint) +# --------------------------------------------------------------------------- + + +@pytest.mark.slow +@has_checkpoint +def test_smoke_2048_full() -> None: + """Full 2048 inference with real checkpoint.""" + engine = CorridorKeyMLXEngine( + checkpoint_path=MLX_CHECKPOINT_PATH, + img_size=2048, + compile=False, + ) + + rng = np.random.default_rng(42) + rgb = rng.integers(0, 256, (2048, 2048, 3), dtype=np.uint8) + mask = rng.integers(0, 256, (2048, 2048), dtype=np.uint8) + + result = engine.process_frame(rgb, mask) + + assert result["alpha"].shape == (2048, 2048) + assert result["alpha"].dtype == np.uint8 + assert result["fg"].shape == (2048, 2048, 3) + assert result["comp"].shape == (2048, 2048, 3) + + for key in ("alpha", "fg", "comp"): + arr = result[key] + assert not np.isnan(arr).any(), f"{key} contains NaN" + assert not np.isinf(arr).any(), f"{key} contains Inf" + + assert result["alpha"].min() != result["alpha"].max(), "alpha is constant" diff --git a/tests/test_greenformer_forward.py b/tests/test_greenformer_forward.py deleted file mode 100644 index b720a0f..0000000 --- a/tests/test_greenformer_forward.py +++ /dev/null @@ -1,77 +0,0 @@ -"""Forward pass shape and smoke tests for GreenFormer (Phase 5). - -Uses random weights — no checkpoint needed. -""" - -from __future__ import annotations - -import mlx.core as mx -import pytest - -from corridorkey_mlx.model.corridorkey import GreenFormer - -IMG_SIZE = 512 -BATCH = 1 - - -@pytest.fixture(scope="module") -def model_and_output() -> tuple[GreenFormer, dict[str, mx.array]]: - """Build model with random weights and run forward pass once.""" - model = GreenFormer(img_size=IMG_SIZE) - x = mx.random.normal((BATCH, IMG_SIZE, IMG_SIZE, 4)) - out = model(x) - # materialize all outputs — mx.eval is MLX lazy graph evaluation, not Python eval - mx.eval(out) - return model, out - - -EXPECTED_SHAPES: dict[str, tuple[int, ...]] = { - "alpha_logits": (BATCH, IMG_SIZE // 4, IMG_SIZE // 4, 1), - "fg_logits": (BATCH, IMG_SIZE // 4, IMG_SIZE // 4, 3), - "alpha_logits_up": (BATCH, IMG_SIZE, IMG_SIZE, 1), - "fg_logits_up": (BATCH, IMG_SIZE, IMG_SIZE, 3), - "alpha_coarse": (BATCH, IMG_SIZE, IMG_SIZE, 1), - "fg_coarse": (BATCH, IMG_SIZE, IMG_SIZE, 3), - "delta_logits": (BATCH, IMG_SIZE, IMG_SIZE, 4), - "alpha_final": (BATCH, IMG_SIZE, IMG_SIZE, 1), - "fg_final": (BATCH, IMG_SIZE, IMG_SIZE, 3), -} - - -@pytest.mark.parametrize("key,expected_shape", EXPECTED_SHAPES.items()) -def test_output_shapes( - key: str, - expected_shape: tuple[int, ...], - model_and_output: tuple[GreenFormer, dict[str, mx.array]], -) -> None: - _, out = model_and_output - assert key in out, f"Missing output key: {key}" - assert out[key].shape == expected_shape, ( - f"{key}: expected {expected_shape}, got {out[key].shape}" - ) - - -def test_all_keys_present( - model_and_output: tuple[GreenFormer, dict[str, mx.array]], -) -> None: - _, out = model_and_output - assert set(out.keys()) == set(EXPECTED_SHAPES.keys()) - - -def test_coarse_probs_in_range( - model_and_output: tuple[GreenFormer, dict[str, mx.array]], -) -> None: - """Sigmoid outputs must be in [0, 1].""" - _, out = model_and_output - for key in ("alpha_coarse", "fg_coarse", "alpha_final", "fg_final"): - arr = out[key] - assert float(mx.min(arr)) >= 0.0, f"{key} has values < 0" - assert float(mx.max(arr)) <= 1.0, f"{key} has values > 1" - - -def test_output_dtype( - model_and_output: tuple[GreenFormer, dict[str, mx.array]], -) -> None: - _, out = model_and_output - for key, arr in out.items(): - assert arr.dtype == mx.float32, f"{key}: expected float32, got {arr.dtype}" diff --git a/tests/test_hiera_stage_parity.py b/tests/test_hiera_stage_parity.py deleted file mode 100644 index be9e0c7..0000000 --- a/tests/test_hiera_stage_parity.py +++ /dev/null @@ -1,79 +0,0 @@ -"""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 — 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 deleted file mode 100644 index 0780d39..0000000 --- a/tests/test_hiera_stage_shapes.py +++ /dev/null @@ -1,46 +0,0 @@ -"""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] diff --git a/tests/test_import.py b/tests/test_import.py deleted file mode 100644 index 298f266..0000000 --- a/tests/test_import.py +++ /dev/null @@ -1,12 +0,0 @@ -"""Smoke test: verify package imports succeed.""" - - -def test_package_imports() -> None: - import corridorkey_mlx - import corridorkey_mlx.convert - import corridorkey_mlx.inference - import corridorkey_mlx.io - import corridorkey_mlx.model - import corridorkey_mlx.utils - - assert corridorkey_mlx.__version__ == "0.1.0" diff --git a/tests/test_model_contract.py b/tests/test_model_contract.py new file mode 100644 index 0000000..1b10002 --- /dev/null +++ b/tests/test_model_contract.py @@ -0,0 +1,194 @@ +"""Model contract tests: imports, shapes, outputs, determinism. + +Consolidated from: test_import.py, test_greenformer_forward.py, +test_hiera_stage_shapes.py, test_end_to_end_smoke.py, test_smoke_2048.py (wiring). + +No checkpoint needed — uses random weights only. +""" + +from __future__ import annotations + +import mlx.core as mx +import numpy as np +import pytest + +from corridorkey_mlx.model.corridorkey import GreenFormer +from corridorkey_mlx.model.hiera import HieraBackbone, HieraPatchEmbed + +from .conftest import IMG_SIZE, SMALL_IMG_SIZE + +# --------------------------------------------------------------------------- +# Package imports +# --------------------------------------------------------------------------- + + +def test_package_imports() -> None: + import corridorkey_mlx + import corridorkey_mlx.convert # noqa: F401 + import corridorkey_mlx.inference # noqa: F401 + import corridorkey_mlx.io # noqa: F401 + import corridorkey_mlx.model # noqa: F401 + import corridorkey_mlx.utils # noqa: F401 + + assert corridorkey_mlx.__version__ == "0.1.0" + + +# --------------------------------------------------------------------------- +# GreenFormer output contract (random weights) +# --------------------------------------------------------------------------- + +EXPECTED_SHAPES: dict[str, tuple[int, ...]] = { + "alpha_logits": (1, IMG_SIZE // 4, IMG_SIZE // 4, 1), + "fg_logits": (1, IMG_SIZE // 4, IMG_SIZE // 4, 3), + "alpha_logits_up": (1, IMG_SIZE, IMG_SIZE, 1), + "fg_logits_up": (1, IMG_SIZE, IMG_SIZE, 3), + "alpha_coarse": (1, IMG_SIZE, IMG_SIZE, 1), + "fg_coarse": (1, IMG_SIZE, IMG_SIZE, 3), + "delta_logits": (1, IMG_SIZE, IMG_SIZE, 4), + "alpha_final": (1, IMG_SIZE, IMG_SIZE, 1), + "fg_final": (1, IMG_SIZE, IMG_SIZE, 3), +} + + +@pytest.fixture(scope="module") +def model_output() -> dict[str, mx.array]: + """Forward pass with random weights, materialized once.""" + model = GreenFormer(img_size=IMG_SIZE) + x = mx.random.normal((1, IMG_SIZE, IMG_SIZE, 4)) + out = model(x) + # mx.eval is MLX array materialization, not Python eval() + mx.eval(out) # noqa: S307 + return out + + +def test_all_keys_present(model_output: dict[str, mx.array]) -> None: + assert set(model_output.keys()) == set(EXPECTED_SHAPES.keys()) + + +@pytest.mark.parametrize("key,expected_shape", EXPECTED_SHAPES.items()) +def test_output_shapes( + key: str, + expected_shape: tuple[int, ...], + model_output: dict[str, mx.array], +) -> None: + assert model_output[key].shape == expected_shape, ( + f"{key}: expected {expected_shape}, got {model_output[key].shape}" + ) + + +def test_output_dtype_float32(model_output: dict[str, mx.array]) -> None: + for key, arr in model_output.items(): + assert arr.dtype == mx.float32, f"{key}: expected float32, got {arr.dtype}" + + +def test_sigmoid_outputs_in_range(model_output: dict[str, mx.array]) -> None: + """Post-sigmoid outputs must be in [0, 1].""" + for key in ("alpha_coarse", "fg_coarse", "alpha_final", "fg_final"): + arr = model_output[key] + assert float(mx.min(arr)) >= 0.0, f"{key} has values < 0" + assert float(mx.max(arr)) <= 1.0, f"{key} has values > 1" + + +# --------------------------------------------------------------------------- +# Backbone shape contract +# --------------------------------------------------------------------------- + +BACKBONE_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: + 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 = HieraBackbone(img_size=IMG_SIZE) + x = mx.zeros((1, IMG_SIZE, IMG_SIZE, 4)) + features = backbone(x) + assert len(features) == 4 + + +@pytest.mark.parametrize("stage_idx", range(4)) +def test_backbone_feature_shapes(stage_idx: int) -> None: + backbone = HieraBackbone(img_size=IMG_SIZE) + x = mx.zeros((1, IMG_SIZE, IMG_SIZE, 4)) + features = backbone(x) + assert features[stage_idx].shape == BACKBONE_SHAPES[stage_idx] + + +# --------------------------------------------------------------------------- +# Full pipeline roundtrip (preprocess -> model -> postprocess) +# --------------------------------------------------------------------------- + + +def test_pipeline_roundtrip() -> None: + """numpy -> preprocess -> model -> postprocess -> uint8.""" + from corridorkey_mlx.io.image import postprocess_alpha, postprocess_foreground, preprocess + + model = GreenFormer(img_size=SMALL_IMG_SIZE) + rgb = np.random.rand(SMALL_IMG_SIZE, SMALL_IMG_SIZE, 3).astype(np.float32) + alpha_hint = np.random.rand(SMALL_IMG_SIZE, SMALL_IMG_SIZE, 1).astype(np.float32) + + x = preprocess(rgb, alpha_hint) + assert x.shape == (1, SMALL_IMG_SIZE, SMALL_IMG_SIZE, 4) + + out = model(x) + mx.eval(out) # noqa: S307 + + alpha = postprocess_alpha(out["alpha_final"]) + fg = postprocess_foreground(out["fg_final"]) + + assert alpha.shape == (SMALL_IMG_SIZE, SMALL_IMG_SIZE) + assert alpha.dtype == np.uint8 + assert fg.shape == (SMALL_IMG_SIZE, SMALL_IMG_SIZE, 3) + assert fg.dtype == np.uint8 + + +# --------------------------------------------------------------------------- +# Determinism +# --------------------------------------------------------------------------- + + +def test_deterministic_output() -> None: + """Same input -> same output.""" + model = GreenFormer(img_size=SMALL_IMG_SIZE) + mx.random.seed(123) + x = mx.random.normal((1, SMALL_IMG_SIZE, SMALL_IMG_SIZE, 4)) + mx.eval(x) # noqa: S307 + + out1 = model(x) + mx.eval(out1) # noqa: S307 + out2 = model(x) + mx.eval(out2) # noqa: S307 + + for key in out1: + np.testing.assert_array_equal( + np.array(out1[key]), np.array(out2[key]), err_msg=f"{key} not deterministic" + ) + + +# --------------------------------------------------------------------------- +# No NaN/Inf (wiring check) +# --------------------------------------------------------------------------- + + +def test_no_nan_inf() -> None: + """Random-weight forward produces finite outputs.""" + model = GreenFormer(img_size=SMALL_IMG_SIZE) + x = mx.random.normal((1, SMALL_IMG_SIZE, SMALL_IMG_SIZE, 4)) + mx.eval(x) # noqa: S307 + out = model(x) + mx.eval(out) # noqa: S307 + + for key in ("alpha_final", "fg_final"): + arr = np.array(out[key]) + assert not np.isnan(arr).any(), f"{key} contains NaN" + assert not np.isinf(arr).any(), f"{key} contains Inf" diff --git a/tests/test_parity.py b/tests/test_parity.py new file mode 100644 index 0000000..024e495 --- /dev/null +++ b/tests/test_parity.py @@ -0,0 +1,140 @@ +"""Consolidated golden reference parity tests. + +Merged from: test_end_to_end_parity.py, test_decoder_parity.py, +test_refiner_parity.py, test_hiera_stage_parity.py, test_reference_fixtures.py. + +Loads full model with checkpoint, runs forward on golden input, +compares all outputs against PyTorch golden.npz references. +""" + +from __future__ import annotations + +import mlx.core as mx +import numpy as np +import pytest + +from corridorkey_mlx.model.corridorkey import GreenFormer +from corridorkey_mlx.utils.layout import nchw_to_nhwc_np, nhwc_to_nchw_np + +from .conftest import ( + GOLDEN_PATH, + IMG_SIZE, + MLX_CHECKPOINT_PATH, + PARITY_TOL_BACKBONE, +) + +NUM_BACKBONE_STAGES = 4 + +# Per-output tolerances — backbone drift cascades through decoder/refiner. +OUTPUT_TOLERANCES: dict[str, float] = { + "alpha_logits": 1e-3, + "fg_logits": 1e-3, + "alpha_logits_up": 1e-3, + "fg_logits_up": 1e-3, + "alpha_coarse": 1e-4, + "fg_coarse": 1e-4, + "delta_logits": 2e-3, + "alpha_final": 2e-4, + "fg_final": 2e-4, +} + + +def _skip_if_missing() -> None: + if not GOLDEN_PATH.exists(): + pytest.skip("golden.npz not found") + if not MLX_CHECKPOINT_PATH.exists(): + pytest.skip("MLX checkpoint not found") + + +@pytest.fixture(scope="module") +def model_outputs_and_fixtures() -> tuple[dict[str, mx.array], dict[str, np.ndarray]]: + """Load model, run forward with golden input, return (mlx_outputs, fixtures).""" + _skip_if_missing() + + fixtures = dict(np.load(GOLDEN_PATH)) + + # NCHW -> NHWC for MLX + input_nhwc = mx.array(nchw_to_nhwc_np(fixtures["input"])) + + model = GreenFormer(img_size=IMG_SIZE) + model.load_checkpoint(MLX_CHECKPOINT_PATH) + + outputs = model(input_nhwc) + # mx.eval is MLX array materialization, not Python eval() + mx.eval(outputs) # noqa: S307 + + return outputs, fixtures + + +# --------------------------------------------------------------------------- +# Backbone stage parity (4 stages) +# --------------------------------------------------------------------------- + + +@pytest.fixture(scope="module") +def backbone_features() -> tuple[list[mx.array], dict[str, np.ndarray]]: + """Run backbone alone to get intermediate stage features.""" + _skip_if_missing() + + fixtures = dict(np.load(GOLDEN_PATH)) + input_nhwc = mx.array(nchw_to_nhwc_np(fixtures["input"])) + + from corridorkey_mlx.model.hiera import HieraBackbone + + backbone = HieraBackbone(img_size=IMG_SIZE) + backbone.load_checkpoint(MLX_CHECKPOINT_PATH) + + features = backbone(input_nhwc) + mx.eval(features) # noqa: S307 + + return features, fixtures + + +@pytest.mark.parametrize("stage_idx", range(NUM_BACKBONE_STAGES)) +def test_backbone_stage_parity( + stage_idx: int, + backbone_features: tuple[list[mx.array], dict[str, np.ndarray]], +) -> None: + """MLX backbone stage output matches PyTorch within tolerance.""" + features, fixtures = backbone_features + + expected_nchw = fixtures[f"encoder_feature_{stage_idx}"] + result_nchw = nhwc_to_nchw_np(np.array(features[stage_idx])) + + assert result_nchw.shape == expected_nchw.shape, ( + f"Stage {stage_idx} shape mismatch: {result_nchw.shape} vs {expected_nchw.shape}" + ) + + max_abs_err = float(np.max(np.abs(result_nchw - expected_nchw))) + assert max_abs_err < PARITY_TOL_BACKBONE, ( + f"Stage {stage_idx} max_abs={max_abs_err:.2e} > {PARITY_TOL_BACKBONE}" + ) + + +# --------------------------------------------------------------------------- +# E2E output parity (decoder + refiner outputs) +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("key", list(OUTPUT_TOLERANCES.keys())) +def test_output_parity( + key: str, + model_outputs_and_fixtures: tuple[dict[str, mx.array], dict[str, np.ndarray]], +) -> None: + """MLX output matches PyTorch golden reference within tolerance.""" + outputs, fixtures = model_outputs_and_fixtures + + assert key in outputs, f"Missing MLX output key: {key}" + assert key in fixtures, f"Missing fixture key: {key}" + + mlx_nchw = nhwc_to_nchw_np(np.array(outputs[key])) + expected_nchw = fixtures[key] + + assert mlx_nchw.shape == expected_nchw.shape, ( + f"{key} shape mismatch: {mlx_nchw.shape} vs {expected_nchw.shape}" + ) + + max_abs = float(np.max(np.abs(mlx_nchw - expected_nchw))) + tol = OUTPUT_TOLERANCES[key] + + assert max_abs < tol, f"{key}: max_abs={max_abs:.2e} > tol={tol:.1e}" diff --git a/tests/test_reference_fixtures.py b/tests/test_reference_fixtures.py deleted file mode 100644 index 5ce3a42..0000000 --- a/tests/test_reference_fixtures.py +++ /dev/null @@ -1,130 +0,0 @@ -"""Parity tests: PyTorch reference fixtures (Phase 1). - -These tests validate that dump_pytorch_reference.py produces -fixtures with expected shapes and dtypes. -""" - -from __future__ import annotations - -from pathlib import Path - -import numpy as np -import pytest - -FIXTURE_PATH = Path("reference/fixtures/golden.npz") - -IMG_SIZE = 512 -BATCH = 1 - -# Expected shapes for each tensor in the fixture file -EXPECTED_SHAPES: dict[str, tuple[int, ...]] = { - "input": (BATCH, 4, IMG_SIZE, IMG_SIZE), - # Backbone features at strides 4, 8, 16, 32 - "encoder_feature_0": (BATCH, 112, IMG_SIZE // 4, IMG_SIZE // 4), - "encoder_feature_1": (BATCH, 224, IMG_SIZE // 8, IMG_SIZE // 8), - "encoder_feature_2": (BATCH, 448, IMG_SIZE // 16, IMG_SIZE // 16), - "encoder_feature_3": (BATCH, 896, IMG_SIZE // 32, IMG_SIZE // 32), - # Decoder outputs at H/4 resolution - "alpha_logits": (BATCH, 1, IMG_SIZE // 4, IMG_SIZE // 4), - "fg_logits": (BATCH, 3, IMG_SIZE // 4, IMG_SIZE // 4), - # Upsampled to full resolution - "alpha_logits_up": (BATCH, 1, IMG_SIZE, IMG_SIZE), - "fg_logits_up": (BATCH, 3, IMG_SIZE, IMG_SIZE), - # Coarse predictions (after sigmoid) - "alpha_coarse": (BATCH, 1, IMG_SIZE, IMG_SIZE), - "fg_coarse": (BATCH, 3, IMG_SIZE, IMG_SIZE), - # Refiner - "delta_logits": (BATCH, 4, IMG_SIZE, IMG_SIZE), - # Final outputs - "alpha_final": (BATCH, 1, IMG_SIZE, IMG_SIZE), - "fg_final": (BATCH, 3, IMG_SIZE, IMG_SIZE), -} - - -@pytest.fixture() -def fixtures() -> dict[str, np.ndarray]: - if not FIXTURE_PATH.exists(): - pytest.skip(f"Fixture file not found: {FIXTURE_PATH}") - return dict(np.load(FIXTURE_PATH)) - - -class TestFixtureCompleteness: - """All expected tensors exist in the fixture file.""" - - def test_all_keys_present(self, fixtures: dict[str, np.ndarray]) -> None: - missing = set(EXPECTED_SHAPES.keys()) - set(fixtures.keys()) - assert not missing, f"Missing fixture keys: {missing}" - - def test_no_extra_keys(self, fixtures: dict[str, np.ndarray]) -> None: - extra = set(fixtures.keys()) - set(EXPECTED_SHAPES.keys()) - assert not extra, f"Unexpected fixture keys: {extra}" - - -class TestBackboneFeaturesShape: - """4 feature maps from Hiera backbone have expected shapes.""" - - @pytest.mark.parametrize("idx", range(4)) - def test_feature_shape(self, fixtures: dict[str, np.ndarray], idx: int) -> None: - key = f"encoder_feature_{idx}" - assert fixtures[key].shape == EXPECTED_SHAPES[key] - - @pytest.mark.parametrize("idx", range(4)) - def test_feature_dtype(self, fixtures: dict[str, np.ndarray], idx: int) -> None: - key = f"encoder_feature_{idx}" - assert fixtures[key].dtype == np.float32 - - -class TestCoarsePredictionsShape: - """Alpha (1ch) and foreground (3ch) coarse logits/probs exist with correct shapes.""" - - @pytest.mark.parametrize( - "key", - [ - "alpha_logits", - "fg_logits", - "alpha_logits_up", - "fg_logits_up", - "alpha_coarse", - "fg_coarse", - ], - ) - def test_shape(self, fixtures: dict[str, np.ndarray], key: str) -> None: - assert fixtures[key].shape == EXPECTED_SHAPES[key] - - @pytest.mark.parametrize( - "key", - [ - "alpha_logits", - "fg_logits", - "alpha_logits_up", - "fg_logits_up", - "alpha_coarse", - "fg_coarse", - ], - ) - def test_dtype(self, fixtures: dict[str, np.ndarray], key: str) -> None: - assert fixtures[key].dtype == np.float32 - - @pytest.mark.parametrize("key", ["alpha_coarse", "fg_coarse"]) - def test_sigmoid_range(self, fixtures: dict[str, np.ndarray], key: str) -> None: - """Coarse predictions (post-sigmoid) should be in [0, 1].""" - assert fixtures[key].min() >= 0.0 - assert fixtures[key].max() <= 1.0 - - -class TestRefinerOutputsShape: - """Delta logits and final alpha/fg have expected shapes.""" - - @pytest.mark.parametrize("key", ["delta_logits", "alpha_final", "fg_final"]) - def test_shape(self, fixtures: dict[str, np.ndarray], key: str) -> None: - assert fixtures[key].shape == EXPECTED_SHAPES[key] - - @pytest.mark.parametrize("key", ["delta_logits", "alpha_final", "fg_final"]) - def test_dtype(self, fixtures: dict[str, np.ndarray], key: str) -> None: - assert fixtures[key].dtype == np.float32 - - @pytest.mark.parametrize("key", ["alpha_final", "fg_final"]) - def test_final_sigmoid_range(self, fixtures: dict[str, np.ndarray], key: str) -> None: - """Final predictions (post-sigmoid) should be in [0, 1].""" - assert fixtures[key].min() >= 0.0 - assert fixtures[key].max() <= 1.0 diff --git a/tests/test_refiner_parity.py b/tests/test_refiner_parity.py deleted file mode 100644 index 7289036..0000000 --- a/tests/test_refiner_parity.py +++ /dev/null @@ -1,154 +0,0 @@ -"""Parity tests: MLX refiner vs PyTorch reference (Phase 2). - -Uses saved coarse predictions + RGB -> runs MLX refiner -> -compares against saved PyTorch delta logits and final outputs. -""" - -from __future__ import annotations - -from pathlib import Path - -import mlx.core as mx -import numpy as np -import pytest - -from corridorkey_mlx.model.refiner import CNNRefinerModule -from corridorkey_mlx.utils.layout import conv_weight_pt_to_mlx, nchw_to_nhwc_np, nhwc_to_nchw_np - -FIXTURE_PATH = Path("reference/fixtures/golden.npz") -WEIGHTS_PATH = Path("reference/fixtures/golden_weights.npz") - -MAX_ABS_TOL = 1e-4 - - -def _skip_if_missing() -> None: - if not FIXTURE_PATH.exists() or not WEIGHTS_PATH.exists(): - pytest.skip("Fixture files not found — run dump_pytorch_reference.py first") - - -def _load_refiner_weights(weights: dict[str, np.ndarray]) -> list[tuple[str, mx.array]]: - """Convert PyTorch refiner state_dict to MLX weight list. - - Handles: - - Conv weight transpose (O,I,H,W) -> (O,H,W,I) - - Key mapping from PyTorch Sequential indices to named attributes - """ - prefix = "refiner." - weight_list: list[tuple[str, mx.array]] = [] - - # Map PyTorch stem Sequential keys to our named attributes - stem_key_map = { - "stem.0.weight": "stem_conv.weight", - "stem.0.bias": "stem_conv.bias", - "stem.1.weight": "stem_gn.weight", - "stem.1.bias": "stem_gn.bias", - } - - for pt_key, value in weights.items(): - if not pt_key.startswith(prefix): - continue - mlx_key = pt_key[len(prefix) :] - - # Remap stem keys - if mlx_key in stem_key_map: - mlx_key = stem_key_map[mlx_key] - - # Conv2d weights: (O,I,H,W) -> (O,H,W,I) - if mlx_key.endswith(".weight") and _is_conv_weight(value): - value = conv_weight_pt_to_mlx(value) - - weight_list.append((mlx_key, mx.array(value))) - - return weight_list - - -def _is_conv_weight(value: np.ndarray) -> bool: - """Check if a weight tensor is a conv weight (4D with spatial dims).""" - return value.ndim == 4 - - -def _build_and_load_refiner(weights: dict[str, np.ndarray]) -> CNNRefinerModule: - """Build an MLX CNNRefinerModule and load converted weights.""" - refiner = CNNRefinerModule() - weight_list = _load_refiner_weights(weights) - refiner.load_weights(weight_list) - refiner.eval() - mx.eval(refiner.parameters()) # noqa: S307 — mx.eval is MLX's compute trigger - return refiner - - -def test_refiner_delta_parity() -> None: - """MLX refiner delta logits match PyTorch within tolerance.""" - _skip_if_missing() - - fixtures = dict(np.load(FIXTURE_PATH)) - weights = dict(np.load(WEIGHTS_PATH)) - - # RGB: first 3 channels of input - rgb_nchw = fixtures["input"][:, :3] # (1, 3, H, W) - rgb_nhwc = mx.array(nchw_to_nhwc_np(rgb_nchw)) - - # Coarse predictions: alpha_coarse (1ch) + fg_coarse (3ch) = 4ch - alpha_coarse_nhwc = nchw_to_nhwc_np(fixtures["alpha_coarse"]) - fg_coarse_nhwc = nchw_to_nhwc_np(fixtures["fg_coarse"]) - coarse_pred = mx.array(np.concatenate([alpha_coarse_nhwc, fg_coarse_nhwc], axis=-1)) - - refiner = _build_and_load_refiner(weights) - result_nhwc = refiner(rgb_nhwc, coarse_pred) - mx.eval(result_nhwc) # noqa: S307 - - result_nchw = nhwc_to_nchw_np(np.array(result_nhwc)) - expected = fixtures["delta_logits"] - - max_abs_err = float(np.max(np.abs(result_nchw - expected))) - mean_abs_err = float(np.mean(np.abs(result_nchw - expected))) - print(f"\nRefiner delta parity — max_abs: {max_abs_err:.6e}, mean_abs: {mean_abs_err:.6e}") - - assert result_nchw.shape == expected.shape - assert max_abs_err < MAX_ABS_TOL, ( - f"Max abs error {max_abs_err:.6e} exceeds tolerance {MAX_ABS_TOL}" - ) - - -def test_refiner_final_output_parity() -> None: - """MLX final alpha and fg match PyTorch within tolerance. - - Tests the full residual + sigmoid path using MLX refiner output. - """ - _skip_if_missing() - - fixtures = dict(np.load(FIXTURE_PATH)) - weights = dict(np.load(WEIGHTS_PATH)) - - # Run refiner - rgb_nchw = fixtures["input"][:, :3] - rgb_nhwc = mx.array(nchw_to_nhwc_np(rgb_nchw)) - alpha_coarse_nhwc = nchw_to_nhwc_np(fixtures["alpha_coarse"]) - fg_coarse_nhwc = nchw_to_nhwc_np(fixtures["fg_coarse"]) - coarse_pred = mx.array(np.concatenate([alpha_coarse_nhwc, fg_coarse_nhwc], axis=-1)) - - refiner = _build_and_load_refiner(weights) - delta_nhwc = refiner(rgb_nhwc, coarse_pred) - mx.eval(delta_nhwc) # noqa: S307 - - # Apply residual in logit space + sigmoid (matching PyTorch forward pass) - alpha_logits_up = mx.array(nchw_to_nhwc_np(fixtures["alpha_logits_up"])) - fg_logits_up = mx.array(nchw_to_nhwc_np(fixtures["fg_logits_up"])) - - alpha_final = mx.sigmoid(alpha_logits_up + delta_nhwc[..., 0:1]) - fg_final = mx.sigmoid(fg_logits_up + delta_nhwc[..., 1:4]) - mx.eval(alpha_final, fg_final) # noqa: S307 - - # Compare - for name, result, expected_key in [ - ("alpha_final", alpha_final, "alpha_final"), - ("fg_final", fg_final, "fg_final"), - ]: - result_nchw = nhwc_to_nchw_np(np.array(result)) - expected = fixtures[expected_key] - max_abs_err = float(np.max(np.abs(result_nchw - expected))) - mean_abs_err = float(np.mean(np.abs(result_nchw - expected))) - print(f"\n{name} parity — max_abs: {max_abs_err:.6e}, mean_abs: {mean_abs_err:.6e}") - assert max_abs_err < MAX_ABS_TOL, ( - f"{name}: max abs error {max_abs_err:.6e} exceeds tolerance {MAX_ABS_TOL}" - ) diff --git a/tests/test_smoke_2048.py b/tests/test_smoke_2048.py deleted file mode 100644 index 67d341c..0000000 --- a/tests/test_smoke_2048.py +++ /dev/null @@ -1,78 +0,0 @@ -"""2048 smoke tests — execution check at native resolution. - -test_smoke_2048_wiring: lightweight, no checkpoint, runs in normal suite. -test_smoke_2048_full: loads real checkpoint at 2048, marked slow + skipif. -""" - -from __future__ import annotations - -from pathlib import Path - -import mlx.core as mx -import numpy as np -import pytest - -from corridorkey_mlx.model.corridorkey import GreenFormer - -CHECKPOINT = Path("checkpoints/corridorkey_mlx.safetensors") -HAS_CHECKPOINT = CHECKPOINT.exists() - -WIRING_IMG_SIZE = 256 # small enough to run fast with random weights - - -def test_smoke_2048_wiring() -> None: - """GreenFormer constructs at 2048 and produces correct shapes (random weights).""" - model = GreenFormer(img_size=WIRING_IMG_SIZE) - x = mx.random.normal((1, WIRING_IMG_SIZE, WIRING_IMG_SIZE, 4)) - # mx.eval materializes lazy MLX arrays (not Python eval) - mx.eval(x) # noqa: S307 - - out = model(x) - mx.eval(out) # noqa: S307 - - assert out["alpha_final"].shape == (1, WIRING_IMG_SIZE, WIRING_IMG_SIZE, 1) - assert out["fg_final"].shape == (1, WIRING_IMG_SIZE, WIRING_IMG_SIZE, 3) - - # No NaN/Inf - alpha = np.array(out["alpha_final"]) - fg = np.array(out["fg_final"]) - assert not np.isnan(alpha).any(), "alpha_final contains NaN" - assert not np.isnan(fg).any(), "fg_final contains NaN" - assert not np.isinf(alpha).any(), "alpha_final contains Inf" - assert not np.isinf(fg).any(), "fg_final contains Inf" - - -@pytest.mark.slow -@pytest.mark.skipif(not HAS_CHECKPOINT, reason="Checkpoint not available") -def test_smoke_2048_full() -> None: - """Full 2048 inference with real checkpoint via engine.""" - from corridorkey_mlx import CorridorKeyMLXEngine - - engine = CorridorKeyMLXEngine( - checkpoint_path=CHECKPOINT, - img_size=2048, - compile=False, # skip compile overhead in test - ) - - # Use synthetic inputs at 2048 - rng = np.random.default_rng(42) - rgb = rng.integers(0, 256, (2048, 2048, 3), dtype=np.uint8) - mask = rng.integers(0, 256, (2048, 2048), dtype=np.uint8) - - result = engine.process_frame(rgb, mask) - - # Shape checks - assert result["alpha"].shape == (2048, 2048) - assert result["alpha"].dtype == np.uint8 - assert result["fg"].shape == (2048, 2048, 3) - assert result["fg"].dtype == np.uint8 - assert result["comp"].shape == (2048, 2048, 3) - - # No NaN/Inf - for key in ("alpha", "fg", "comp"): - arr = result[key] - assert not np.isnan(arr).any(), f"{key} contains NaN" - assert not np.isinf(arr).any(), f"{key} contains Inf" - - # Alpha shouldn't be degenerate - assert result["alpha"].min() != result["alpha"].max(), "alpha is constant — suspicious" diff --git a/tests/test_tiling_consistency.py b/tests/test_tiling.py similarity index 100% rename from tests/test_tiling_consistency.py rename to tests/test_tiling.py diff --git a/tests/test_weights.py b/tests/test_weights.py index 9674e78..d23d998 100644 --- a/tests/test_weights.py +++ b/tests/test_weights.py @@ -1,15 +1,9 @@ -"""Unit tests for weights download module. - -Tests config, caching, checksum verification, and CLI arg parsing -without requiring real network access. -""" +"""Tests for weights download module: config defaults, checksum, hash parsing.""" from __future__ import annotations import hashlib -import os from typing import TYPE_CHECKING -from unittest.mock import patch from corridorkey_mlx.weights import ( DEFAULT_ASSET_NAME, @@ -22,7 +16,6 @@ from corridorkey_mlx.weights import ( github_owner_repo, verify_sha256, ) -from corridorkey_mlx.weights_cli import build_parser if TYPE_CHECKING: from pathlib import Path @@ -37,26 +30,12 @@ class TestConfig: assert owner == DEFAULT_GITHUB_OWNER assert repo == DEFAULT_GITHUB_REPO - def test_owner_repo_env_override(self) -> None: - with patch.dict(os.environ, {"CORRIDORKEY_MLX_WEIGHTS_REPO": "other/repo"}): - owner, repo = github_owner_repo() - assert owner == "other" - assert repo == "repo" - def test_default_asset_name(self) -> None: assert default_asset_name() == DEFAULT_ASSET_NAME - def test_asset_name_env_override(self) -> None: - with patch.dict(os.environ, {"CORRIDORKEY_MLX_WEIGHTS_ASSET": "custom.st"}): - assert default_asset_name() == "custom.st" - def test_default_tag(self) -> None: assert default_tag() == "latest" - def test_tag_env_override(self) -> None: - with patch.dict(os.environ, {"CORRIDORKEY_MLX_WEIGHTS_TAG": "v2.0"}): - assert default_tag() == "v2.0" - # -- Cache directory --------------------------------------------------------- @@ -91,45 +70,14 @@ class TestChecksum: def test_parse_hash_garbage(self) -> None: assert _parse_hash_line("not a hash", "file.bin") is None - def test_verify_sha256(self, tmp_path: Path) -> None: + def test_verify_sha256_valid(self, tmp_path: Path) -> None: content = b"hello world" expected = hashlib.sha256(content).hexdigest() p = tmp_path / "test.bin" p.write_bytes(content) - assert verify_sha256(p, expected) is True + + def test_verify_sha256_invalid(self, tmp_path: Path) -> None: + p = tmp_path / "test.bin" + p.write_bytes(b"hello world") assert verify_sha256(p, "0" * 64) is False - - -# -- CLI arg parsing --------------------------------------------------------- - - -class TestCLI: - def test_download_defaults(self) -> None: - parser = build_parser() - args = parser.parse_args(["download"]) - assert args.command == "download" - assert args.tag is None - assert args.asset_name is None - assert args.out is None - assert args.force is False - assert args.no_verify is False - assert args.print_path is False - - def test_download_all_flags(self) -> None: - parser = build_parser() - args = parser.parse_args([ - "download", - "--tag", "v1.0", - "--asset", "custom.st", - "--out", "/tmp/w", - "--force", - "--no-verify", - "--print-path", - ]) - assert args.tag == "v1.0" - assert args.asset_name == "custom.st" - assert args.out == "/tmp/w" - assert args.force is True - assert args.no_verify is True - assert args.print_path is True