Fix QuantizedLinear for non-affine quantization modes (mxfp4)

[Nicolas-nwb]

Proposed changes

QuantizedLinear fails to load mxfp4-quantized models with .keyNotFound for "biases". Two bugs:

1. Mode not forwarded during quantization — The weight-based init calls MLX.quantized(weight, groupSize:bits:) without passing mode, so weights are always quantized as affine regardless of the mode specified. QuantizedEmbedding already passes mode correctly.

2. Missing biases treated as error for non-affine modes — Module.update(parameters:, verify: .all) calls updateMissing() for the "biases" key, which throws .keyNotFound. Non-affine modes (e.g. .mxfp4) don't produce biases, so this key is legitimately absent.

Fix

• Pass mode: mode to MLX.quantized() in QuantizedLinear.init(weight:bias:groupSize:bits:mode:)
• Override updateMissing() to skip "biases" when mode != .affine

Tests

• testQuantizedLinearMxfp4DoesNotCreateAffineBiases — verifies mxfp4 produces nil biases
• testQuantizedLinearMxfp4ParametersRoundTripWithoutBiases — verifies update(parameters:, verify: .all) succeeds without a "biases" key

Checklist

• I have read the CONTRIBUTING document
• I have run pre-commit run --all-files to format my code / installed pre-commit prior to committing changes
• I have added tests that prove my fix is effective or that my feature works
• I have updated the necessary documentation (if needed)

[davidkoski]

I am curious about this -- the init for QuantizedLinear should have left us with a biases = nil and we shouldn't hit this case. I wonder what happened here?

[davidkoski]

Ah, perhaps this was an earlier attempt to fix it?

this is from the test and biases is indeed nil. I think this method should be removed because the missing parameter will be handled correctly -- this is redundant and confuses the issue by suggesting we need special handling.

[davidkoski]

Please look at my comment on updateMissing() -- I think it should be removed, but let me know if you think otherwise.

Nice find and fix!

[Nicolas-nwb]

You're right — I traced through Module.update(parameters:verify:) and when biases is nil, build(value:) wraps Optional<MLXArray>.none as .value(.other(...)), which hits the (.value(.other(_)), .none) → break case in apply. So updateMissing is never actually called for a nil optional parameter. Removed the override and the round-trip test.

The mxfp4 nil-biases test stays to guard the mode: forwarding fix.

Thanks for the thorough review and for maintaining this project!

[davidkoski]

Changes look good, thank you for figuring this out!