Raise a clear error for an invalid quantization mode in nn layers

[ayaangazali]

Proposed changes

mx.quantize validates its mode argument and gives a clear message, but the nn quantized layers hit a bare dict lookup first, so a typo'd mode surfaces as a raw KeyError with no hint about what is wrong or which modes exist:

import mlx.core as mx
import mlx.nn as nn

mx.quantize(mx.random.uniform(shape=(8, 8)), 64, 4, mode="bogus")
# ValueError: [quantize] Invalid quantization mode 'bogus'.

nn.QuantizedLinear(8, 8, mode="bogus")
# KeyError: 'bogus'

Same KeyError for QuantizedEmbedding, QQLinear, from_linear, from_embedding, and nn.quantize(model, mode=...), since all of them call _defaults_for_mode before anything reaches mx.quantize.

This adds the missing check in _defaults_for_mode so those paths raise a ValueError that names the valid modes instead. Added the cases to test_quantize; they fail on main with KeyError: 'xyz' and pass with the change.

Checklist

Put an x in the boxes that apply.

• 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)

(test_nn.py 69/69 and test_quantized.py 32/32 pass locally on CPU)

---

heads up on process: i'm a freshman contributor and i use Claude Code to help me hunt for things like this, but i reproduced every one of those KeyError paths myself, checked that the valid modes list matches mode_defaults, and confirmed the new test actually fails on main before opening this. if you'd rather word the error differently or put the check somewhere else, happy to change it.

[nastya236]

Thank you for adding this. I don't think this kind of change requires a test. Do you mind dropping the test?

[ayaangazali]

Sure thing, test dropped. It is just the validation check now.

[ayaangazali]

Quick note on the red CI here: the failures are test_quantized.TestQuantized.test_gather_qmm_sorted producing nan, which is unrelated to this PR (this diff only turns a KeyError into a ValueError in _defaults_for_mode, and test_mode_error_cases passes in the same run).

The same test is failing on main right now, for example https://github.com/ml-explore/mlx/actions/runs/30585439060. Leaving it alone rather than rerunning since it looks like it would just reproduce. Happy to rebase once main is green if that helps.