laguna: load the public checkpoint layouts (original poolside + repacks)

[pierre427]

The laguna module from #1334 loads checkpoints whose keys already match its module tree, but the public Laguna checkpoints on the Hub use two other layouts, and all of them currently fail load_weights(strict=True):

Original poolside layout (poolside/Laguna-S-2.1-bf16): experts are stored individually rather than stacked for SwitchGLU, and the router correction bias lives under experts.:

ValueError: Received 36143 parameters not in model:
model.layers.1.mlp.experts.0.down_proj.weight, ...

Public repacks (AtomicChat/Laguna-XS-2.1-MLX-8bit, mlx-community/Laguna-XS-2.1-bf16): every tensor is wrapped in a VLM-style language_model. prefix, and the router is stored as a Linear (mlp.gate.proj.weight) — quantized to weight/scales/biases in the 8-bit repack:

ValueError: Received 1634 parameters not in model:
language_model.lm_head.biases, language_model.lm_head.scales, ...

Changes

• Model.sanitize(): strips the language_model. prefix, stacks individually-stored experts into the SwitchGLU layout, and remaps experts.e_score_correction_bias → gate.e_score_correction_bias.
• MoEGate stores the router as an nn.Linear (gate.proj) instead of a bare matrix, so quantized routers (the 8-bit repack ships one) load directly — a bare matrix can't hold a weight/scales/biases triple, and nn.quantize can't target it either, so mlx_lm.convert -q on a Laguna checkpoint currently leaves the router unquantizable. The bare gate.weight layout (the original checkpoint, and any conversion made with the current module) is mapped onto gate.proj.weight by sanitize(), so existing conversions keep loading. Router math is unchanged.

Verification

• Strict lazy load succeeds on all three public checkpoints above (previously: the errors quoted).
• Full CPU materialization + greedy continuation on AtomicChat/Laguna-XS-2.1-MLX-8bit (33B, quantized router): def quicksort(arr): → "\n if len(arr) <=".
• New test_laguna_sanitize exercises the three layouts (native / prefixed / original per-expert) on a tiny config with strict loading and a forward pass; test_all_models still passes.

🤖 Generated with Claude Code