feat(phase2): MLX decoder + refiner with parity tests

- MLP, DecoderHead (SegFormer-style), RefinerBlock, CNNRefinerModule
- All NHWC native, GroupNorm with pytorch_compatible=True
- Layout utils: NCHW<->NHWC, conv weight transpose (OIHW->OHWI)
- Dump script extended to export decoder/refiner weights
- Parity: decoder max_abs ~5e-6, refiner max_abs ~1e-5

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

docs/plans/2026-03-01-feat-corridorkey-mlx-inference-port-plan.md

@@ -109,10 +109,10 @@ RGB + coarse alpha hint (4ch)
 
 **Components to implement:**
 
-- [ ] `MLP` -- feedforward block
+- [x] `MLP` -- feedforward block
-- [ ] `DecoderHead` -- consumes backbone features, produces coarse predictions
+- [x] `DecoderHead` -- consumes backbone features, produces coarse predictions
-- [ ] `RefinerBlock` -- single refiner stage
+- [x] `RefinerBlock` -- single refiner stage
-- [ ] `CNNRefinerModule` -- full refiner consuming 7ch input
+- [x] `CNNRefinerModule` -- full refiner consuming 7ch input
 
 **Files touched:**
 

scripts/dump_pytorch_reference.py

@@ -333,6 +333,29 @@ def dump_fixtures(outputs: dict[str, torch.Tensor | list[torch.Tensor]], output_
     console.print(f"\n[green]Saved fixtures to {out_path}[/green]")
 
 
+WEIGHTS_FILENAME = "golden_weights.npz"
+
+
+def dump_weights(model: GreenFormer, output_dir: Path) -> None:
+    """Save decoder and refiner weights as numpy arrays for Phase 2 parity tests.
+
+    Keys are the PyTorch state_dict keys (e.g. 'alpha_decoder.linear_c1.proj.weight').
+    Conv weights remain in PyTorch format (O,I,H,W) — tests handle conversion.
+    """
+    output_dir.mkdir(parents=True, exist_ok=True)
+    arrays: dict[str, np.ndarray] = {}
+
+    # Extract decoder and refiner state dicts
+    prefixes = ("alpha_decoder.", "fg_decoder.", "refiner.")
+    for key, param in model.state_dict().items():
+        if any(key.startswith(p) for p in prefixes):
+            arrays[key] = param.cpu().numpy()
+
+    out_path = output_dir / WEIGHTS_FILENAME
+    np.savez(out_path, **arrays)
+    console.print(f"[green]Saved {len(arrays)} weight tensors to {out_path}[/green]")
+
+
 def print_shape_report(outputs: dict[str, torch.Tensor | list[torch.Tensor]]) -> None:
     """Print a rich table of tensor names, shapes, and dtypes."""
     table = Table(title="Reference Fixture Shapes")
@@ -419,6 +442,7 @@ def main() -> None:
 
     print_shape_report(outputs)
     dump_fixtures(outputs, args.output_dir)
+    dump_weights(model, args.output_dir)
 
 
 if __name__ == "__main__":

src/corridorkey_mlx/model/decoder.py

@@ -1,5 +1,95 @@
-"""Decoder heads — MLX port (not yet implemented).
+"""Decoder heads — MLX port.
 
 Two heads: alpha (1ch) and foreground (3ch).
-Consume multiscale backbone features, upsample to full resolution.
+Consume multiscale backbone features (NHWC), upsample and fuse to produce predictions.
+
+Architecture mirrors nikopueringer/CorridorKey DecoderHead (SegFormer-style).
 """
+
+from __future__ import annotations
+
+import mlx.core as mx
+import mlx.nn as nn
+
+
+class MLP(nn.Module):
+    """Single linear projection: input_dim -> embed_dim."""
+
+    def __init__(self, input_dim: int, embed_dim: int) -> None:
+        super().__init__()
+        self.proj = nn.Linear(input_dim, embed_dim)
+
+    def __call__(self, x: mx.array) -> mx.array:
+        return self.proj(x)
+
+
+class DecoderHead(nn.Module):
+    """SegFormer-style multiscale feature fusion head (NHWC).
+
+    Takes 4 multi-scale feature maps in NHWC format, projects each to embed_dim,
+    upsamples to the largest spatial resolution, fuses, and classifies.
+    """
+
+    def __init__(
+        self,
+        in_channels: list[int],
+        embed_dim: int,
+        output_dim: int,
+    ) -> None:
+        super().__init__()
+        self.linear_c1 = MLP(in_channels[0], embed_dim)
+        self.linear_c2 = MLP(in_channels[1], embed_dim)
+        self.linear_c3 = MLP(in_channels[2], embed_dim)
+        self.linear_c4 = MLP(in_channels[3], embed_dim)
+
+        fused_channels = embed_dim * len(in_channels)
+        self.linear_fuse = nn.Conv2d(fused_channels, embed_dim, kernel_size=1, bias=False)
+        self.bn = nn.BatchNorm(embed_dim)
+        self.classifier = nn.Conv2d(embed_dim, output_dim, kernel_size=1)
+
+    def __call__(self, features: list[mx.array]) -> mx.array:
+        """Forward pass.
+
+        Args:
+            features: 4 feature maps in NHWC format.
+                      [0]: (B, H/4,  W/4,  C1)
+                      [1]: (B, H/8,  W/8,  C2)
+                      [2]: (B, H/16, W/16, C3)
+                      [3]: (B, H/32, W/32, C4)
+
+        Returns:
+            Logits in NHWC: (B, H/4, W/4, output_dim)
+        """
+        c1, c2, c3, c4 = features
+        target_h, target_w = c1.shape[1], c1.shape[2]  # H/4, W/4
+
+        projected = []
+        for feat, linear in zip(
+            [c1, c2, c3, c4],
+            [self.linear_c1, self.linear_c2, self.linear_c3, self.linear_c4],
+            strict=True,
+        ):
+            b, h, w, _c = feat.shape
+            # NHWC: reshape to (B, H*W, C), project, reshape back
+            x = feat.reshape(b, h * w, _c)
+            x = linear(x)  # (B, H*W, embed_dim)
+            x = x.reshape(b, h, w, -1)  # (B, H, W, embed_dim)
+            # Upsample to target spatial size
+            if h != target_h or w != target_w:
+                scale_h = target_h / h
+                scale_w = target_w / w
+                upsampler = nn.Upsample(
+                    scale_factor=(scale_h, scale_w),
+                    mode="linear",
+                    align_corners=False,
+                )
+                x = upsampler(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)
+        fused = self.linear_fuse(fused)
+        fused = self.bn(fused)
+        fused = nn.relu(fused)
+        # No dropout at inference (eval mode)
+        return self.classifier(fused)

src/corridorkey_mlx/model/refiner.py

@@ -1,5 +1,87 @@
-"""CNN refiner — MLX port (not yet implemented).
+"""CNN refiner — MLX port.
 
-Input: RGB + coarse predictions (7ch total).
+Input: RGB + coarse predictions (7ch total) in NHWC.
-Output: additive delta logits → final sigmoid.
+Output: additive delta logits (4ch) in NHWC.
+
+Architecture mirrors nikopueringer/CorridorKey CNNRefinerModule.
 """
+
+from __future__ import annotations
+
+import mlx.core as mx
+import mlx.nn as nn
+
+REFINER_CHANNELS = 64
+REFINER_GROUPS = 8
+REFINER_SCALE = 10.0
+
+
+class RefinerBlock(nn.Module):
+    """Dilated residual block with GroupNorm (NHWC)."""
+
+    def __init__(self, channels: int, dilation: int) -> None:
+        super().__init__()
+        self.conv1 = nn.Conv2d(
+            channels,
+            channels,
+            kernel_size=3,
+            padding=dilation,
+            dilation=dilation,
+        )
+        self.gn1 = nn.GroupNorm(REFINER_GROUPS, channels, pytorch_compatible=True)
+        self.conv2 = nn.Conv2d(
+            channels,
+            channels,
+            kernel_size=3,
+            padding=dilation,
+            dilation=dilation,
+        )
+        self.gn2 = nn.GroupNorm(REFINER_GROUPS, channels, pytorch_compatible=True)
+
+    def __call__(self, x: mx.array) -> mx.array:
+        residual = x
+        out = nn.relu(self.gn1(self.conv1(x)))
+        out = self.gn2(self.conv2(out))
+        return nn.relu(out + residual)
+
+
+class CNNRefinerModule(nn.Module):
+    """CNN refiner: stem + 4 dilated ResBlocks + 1x1 projection (NHWC).
+
+    Takes RGB (3ch) and coarse predictions (4ch) concatenated along channel dim.
+    Returns delta logits (4ch) scaled by 10x.
+    """
+
+    def __init__(self) -> None:
+        super().__init__()
+        refiner_input_channels = 7  # RGB (3) + coarse_pred (4: alpha + fg)
+        self.stem_conv = nn.Conv2d(
+            refiner_input_channels, REFINER_CHANNELS, kernel_size=3, padding=1
+        )
+        self.stem_gn = nn.GroupNorm(REFINER_GROUPS, REFINER_CHANNELS, pytorch_compatible=True)
+
+        self.res1 = RefinerBlock(REFINER_CHANNELS, dilation=1)
+        self.res2 = RefinerBlock(REFINER_CHANNELS, dilation=2)
+        self.res3 = RefinerBlock(REFINER_CHANNELS, dilation=4)
+        self.res4 = RefinerBlock(REFINER_CHANNELS, dilation=8)
+
+        refiner_output_channels = 4  # delta for alpha (1) + delta for fg (3)
+        self.final = nn.Conv2d(REFINER_CHANNELS, refiner_output_channels, kernel_size=1)
+
+    def __call__(self, rgb: mx.array, coarse_pred: mx.array) -> mx.array:
+        """Forward pass.
+
+        Args:
+            rgb: (B, H, W, 3) — RGB input in NHWC
+            coarse_pred: (B, H, W, 4) — concatenated alpha_coarse + fg_coarse in NHWC
+
+        Returns:
+            Delta logits: (B, H, W, 4) in NHWC, scaled by 10x
+        """
+        x = mx.concatenate([rgb, coarse_pred], axis=-1)  # (B, H, W, 7)
+        x = nn.relu(self.stem_gn(self.stem_conv(x)))
+        x = self.res1(x)
+        x = self.res2(x)
+        x = self.res3(x)
+        x = self.res4(x)
+        return self.final(x) * REFINER_SCALE

src/corridorkey_mlx/utils/layout.py

@@ -1,4 +1,35 @@
-"""Tensor layout conversion utilities (not yet implemented).
+"""Tensor layout conversion utilities.
 
-Centralizes NCHW ↔ NHWC transforms. All layout conversions go through here.
+Centralizes NCHW <-> NHWC transforms and conv weight transposes.
+All layout conversions in the project go through this module.
 """
+
+from __future__ import annotations
+
+import mlx.core as mx
+import numpy as np
+
+
+def nchw_to_nhwc(x: mx.array) -> mx.array:
+    """Convert (N, C, H, W) -> (N, H, W, C)."""
+    return mx.transpose(x, axes=(0, 2, 3, 1))
+
+
+def nhwc_to_nchw(x: mx.array) -> mx.array:
+    """Convert (N, H, W, C) -> (N, C, H, W)."""
+    return mx.transpose(x, axes=(0, 3, 1, 2))
+
+
+def conv_weight_pt_to_mlx(w: np.ndarray) -> np.ndarray:
+    """Transpose PyTorch conv weight (O, I, H, W) -> MLX conv weight (O, H, W, I)."""
+    return np.transpose(w, axes=(0, 2, 3, 1))
+
+
+def nchw_to_nhwc_np(x: np.ndarray) -> np.ndarray:
+    """Convert numpy (N, C, H, W) -> (N, H, W, C)."""
+    return np.transpose(x, axes=(0, 2, 3, 1))
+
+
+def nhwc_to_nchw_np(x: np.ndarray) -> np.ndarray:
+    """Convert numpy (N, H, W, C) -> (N, C, H, W)."""
+    return np.transpose(x, axes=(0, 3, 1, 2))

tests/test_decoder_parity.py

@@ -1,19 +1,150 @@
 """Parity tests: MLX decoder heads vs PyTorch reference (Phase 2).
 
-Uses saved backbone features → runs MLX decoder → compares
+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
 
-pytestmark = pytest.mark.skip(reason="Phase 2: MLX decoder not yet implemented")
+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",
+    )

tests/test_refiner_parity.py

@@ -1,19 +1,154 @@
 """Parity tests: MLX refiner vs PyTorch reference (Phase 2).
 
-Uses saved coarse predictions + RGB → runs MLX refiner →
+Uses saved coarse predictions + RGB -> runs MLX refiner ->
-compares against saved PyTorch final outputs.
+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
 
-pytestmark = pytest.mark.skip(reason="Phase 2: MLX refiner not yet implemented")
+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."""
+    """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}"
+        )