corridorkey-mrp-mlx/tests/test_conversion.py

"""Tests: weight conversion PyTorch → MLX (Phase 3).

Validates key mapping, shape transforms, and round-trip integrity.
"""

from __future__ import annotations

from pathlib import Path
from typing import TYPE_CHECKING

import numpy as np
import pytest

if TYPE_CHECKING:
    from collections import OrderedDict

from corridorkey_mlx.convert.converter import (
    REFINER_STEM_MAP,
    SKIP_SUFFIXES,
    convert_state_dict,
    load_pytorch_checkpoint,
)

CHECKPOINT_PATH = Path("checkpoints/CorridorKey_v1.0.pth")
SAFETENSORS_PATH = Path("checkpoints/corridorkey_mlx.safetensors")


@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


# ---------------------------------------------------------------------------
# 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")

        from safetensors.numpy import load_file

        loaded = load_file(str(SAFETENSORS_PATH))

        _, converted, _ = checkpoint_data

        assert set(loaded.keys()) == set(converted.keys())
        for key in converted:
            np.testing.assert_array_equal(loaded[key], converted[key])