Repository navigation
Raise a clear error for an invalid quantization mode in nn layers - #3914
Conversation
|
Thank you for adding this. I don't think this kind of change requires a test. Do you mind dropping the test? |
4bfc94c to
e8873dd
Compare
|
Sure thing, test dropped. It is just the validation check now. |
|
Quick note on the red CI here: the failures are 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. |
Proposed changes
mx.quantizevalidates itsmodeargument and gives a clear message, but thennquantized layers hit a bare dict lookup first, so a typo'd mode surfaces as a rawKeyErrorwith no hint about what is wrong or which modes exist:Same
KeyErrorforQuantizedEmbedding,QQLinear,from_linear,from_embedding, andnn.quantize(model, mode=...), since all of them call_defaults_for_modebefore anything reachesmx.quantize.This adds the missing check in
_defaults_for_modeso those paths raise aValueErrorthat names the valid modes instead. Added the cases totest_quantize; they fail on main withKeyError: 'xyz'and pass with the change.Checklist
Put an
xin the boxes that apply.pre-commit run --all-filesto format my code / installed pre-commit prior to committing changes(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
KeyErrorpaths myself, checked that the valid modes list matchesmode_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.