fix(models): build quantized Jamba MoE experts through the shared SwitchGLU

[inureyes]

Summary

Jamba declared a family-local SwitchLinear::Quantized variant that nothing constructed, behind an #[allow(dead_code)], so a checkpoint shipping quantized MoE experts loaded silently and built its expert planes from packed uint32 weights on the dense path. Settling the issue's two readings turned up a second defect: the dense path was itself broken, so Jamba's MoE branch could never run for any checkpoint.

The determination

Reading B, established without the 29 GB download. Full evidence is recorded in #974.

Such checkpoints exist, from two publishers. Confirmed from the Hub API and model.safetensors.index.json alone: ncls-p/AI21-Jamba2-Mini-mlx-4Bit and ssdataanalysis/AI21-Jamba2-Mini-mlx-8Bit both declare num_experts: 16 and ship 144 pre-stacked switch_mlp tensors (48 .weight, 48 .scales, 48 .biases) with zero .experts.N keys. The keys match the loader's lookup verbatim. No mlx-community Jamba conversion has num_experts > 1, which is why the gap survived.

The synthetic route confirmed the load gap and refuted part of the issue's prediction. from_weights returns Ok on a genuinely packed uint32 plane, as predicted. But the first forward never reaches mlxcel_core::matmul, so [matmul] Only inexact types are supported is not the message. It aborts two lines earlier:

libc++abi: terminating due to uncaught exception of type std::invalid_argument:
[reshape] Cannot reshape array of size 256 into shape (1,4,1,1).

and it aborts identically on dense f32 expert planes. SwitchGLU::forward expanded x from [b, s, h] to [b, s, 1, h], then SwitchLinear::Regular::forward expanded it again as if it were still 3D. The comment justifying the local enum by its "different forward shape handling" was describing a bug.

What changed

• src/models/jamba.rs: deleted the family-local SwitchLinear and SwitchGLU and routed the family through switch_layers::SwitchGLU (the issue's step 2 preference). That branches on .scales presence so quantized experts reach gather_qmm, bounds the declared group_size / bits pair and the affine/.biases pairing in from_stacked_parts before storing anything derived from them, and accepts the per-expert experts.{idx} layout through stack_individual_experts.
• src/models/jamba.rs: JambaSparseMoeBlock::forward flattens to the shared [tokens, hidden] / [tokens, top_k] contract and combines through moe_weighted_sum. The top-k routing math is unchanged.
• src/models/jamba.rs: removed the hand-rolled per-expert remap in sanitize_weights, which carried .weight alone and dropped .scales / .biases. The shared loader handles both layouts and derives the same expert keys, and it also runs for callers entering at from_weights, which the removed loop did not.
• src/models/jamba.rs: the router moves from load_linear to UnifiedLinear. Real conversions ship it packed with its own .scales / .biases at a width that need not match the top level (ncls-p declares 8-bit routers under a 4-bit default), and load_linear reads .weight and .bias only.
• src/models/jamba.rs: quantization.mode is now read and bounded at load via mlxcel_core::layers::validate_quantization_mode (fix(core): bound the declared quantization mode string, not just group_size and bits #973). Every Jamba loader builds affine planes, so non-affine is refused with the mode named rather than silently reinterpreted and aborted in the kernel.
• src/models/jamba_tests.rs: four regression tests driving JambaModel::from_weights over pre-stacked, per-expert and dense expert layouts.
• src/models/switch_layers.rs, src/lib/mlxcel-core/src/layers.rs: comment-only. Added Jamba to the Used by: lists on the module, stack_individual_experts, moe_weighted_sum and validate_quantization_mode.

Notes for review

#[allow(dead_code)] is gone because the enum it guarded is gone.

On the issue's step 3: the bound lands via SwitchLinear::from_stacked_parts, which already calls mlxcel_core::layers::validate_quantization_params with the offending tensor prefix in the error. A second validate_expert_quantization_params call would be a duplicate guard, so Jamba is not added to that helper's Used by: list. The tests still walk the shared HOSTILE_QUANT_PARAMS table.

The hostile-parameter walk is the discriminator for the bug, not just a bounds check. The dense gather_mm path never reads the declared pair, so before this change "bits": 0 loaded without complaint. A load that now refuses it is a load that took the quantized branch.

Tests assert on the returned Result and never run a forward pass, per the convention PR #972 settled on: gather_qmm crosses the cxx bridge as UniquePtr<MlxArray> rather than Result, so a regression would abort the test binary with SIGABRT instead of failing cleanly.

Test plan

• models/jamba-v0.1-4bit generates byte-identical output to main on the same raw prompt (-p "The capital of France is" -n 20 --no-chat-template). It is the regression control for the dense path: num_experts: 1 means it never enters the MoE branch.
• cargo fmt --all -- --check
• cargo clippy --release --all-targets --features metal,accelerate -- -D warnings
• cargo clippy -p mlxcel-core --release --all-targets --features metal,accelerate -- -D warnings
• cargo test --release --features metal,accelerate --no-fail-fast: lib 4849 passed, 0 failed. The only two failures are granite4_vision_parity::text_only_forward_produces_finite_logits (fix: Resolve flaky text_only_forward_produces_finite_logits in granite4_vision and hunyuan_vl parity tests #997) and ernie4_5_moe_vl_parity::text_only_forward_produces_finite_logits, both the same concurrent-large-model-load flake; both pass in isolation, verified.
• cargo test --release -p mlxcel-core --features metal,accelerate --no-fail-fast: 1409 passed, 1 failed, the pre-existing fused_moe_parity_tests::fused_moe_geglu_kernel_matches_references_gemma4_shape (fix(tests): fused MoE GeGLU parity threshold is tighter than the gather_qmm reference's own jitter #964).

Not verifiable here: end-to-end generation on a real quantized Jamba MoE checkpoint. The smallest is roughly 29 GB and every AI21-Jamba2-3B MLX conversion is num_experts: 1, so no small MoE Jamba exists to test against. The acceptance criterion asking for an mlx-lm reference comparison on such a checkpoint is left open.

Closes #974