corridorkey-mrp-mlx/prompts/phase4-hiera-backbone.md

Work only on the Hiera backbone port and its stage-by-stage parity.

Goal: Implement an MLX Hiera backbone that reproduces the feature contract expected by CorridorKey’s downstream heads.

Context:

• The original PyTorch model uses a timm backbone created with features_only=True.
• The downstream model expects a pyramid of 4 multiscale feature maps.
• The first input projection is patched from 3 input channels to 4 input channels for RGB + alpha hint.
• This phase is about reproducing backbone behavior and feature outputs, not full end-to-end inference.

Primary deliverables:

• src/corridorkey_mlx/model/hiera.py
• tests/test_hiera_stage_shapes.py
• tests/test_hiera_stage_parity.py
• updates to converter mapping files as needed
• README notes describing the backbone parity status and any known gaps

Requirements:

1. Inspect the original CorridorKey backbone usage and the underlying Hiera/timm feature contract before editing.
2. Treat the PyTorch reference harness as the source of truth for:
   • number of feature stages
   • output order
   • channel counts
   • spatial reductions
3. Implement the smallest MLX Hiera subset necessary to produce the correct 4 feature maps.
4. Preserve the patched 4-channel patch-embed behavior exactly:
   • do not invent a new initialization rule
   • match the PyTorch conversion semantics for the extra alpha-hint channel
5. Keep tensor layout handling explicit and centralized:
   • use one canonical boundary between PyTorch-style NCHW fixtures and MLX NHWC internals
   • do not scatter transpose logic across the backbone code
6. Add narrow parity checks for:
   • patch embed output
   • each stage output
   • final list of multiscale features
7. If exact numerical parity is not immediately possible, first achieve:
   • correct stage count
   • correct stage order
   • correct shapes
   • correct dtype flow then localize the first divergent block
8. Keep the implementation modular:
   • patch embed
   • stage/block definition
   • downsampling / pooling transitions
   • feature collection
9. Update the checkpoint conversion path only as needed to support the backbone weights for completed modules.

Diagnostics to produce:

• stage name
• source tensor shape
• destination tensor shape
• layout assumption
• max abs error
• mean abs error
• whether the mismatch begins before or after a stage transition

Do not:

• wire the full GreenFormer end-to-end yet
• rewrite decoder/refiner code unless required by an interface mismatch
• optimize with compile() yet
• add training code
• silently accept stage reordering or “close enough” feature contracts

Working style:

• Explore first.
• Summarize the exact Hiera feature contract you believe CorridorKey expects before editing.
• After edits, run the narrowest backbone-only tests first.
• When parity fails, stop at the first divergent stage and explain the likely cause before changing more code.

Definition of done:

• MLX backbone returns 4 multiscale feature maps in the correct order
• shape tests pass
• stage-level parity tests exist and run
• the 4-channel patch embed path is implemented and documented
• converter mappings exist for completed backbone weights