E.Huang
Back to projects

mlx-audio

Fixed Canary-1B-v2 checkpoint loading in mlx-audio, with regression tests for full-precision and quantized models. Merged upstream.

PythonMLXASRApple SiliconPyTest

I contributed a fix to mlx-audio’s loader for NVIDIA’s Canary-1B-v2 speech-recognition model. It supported the NeMo-converted checkpoint layout but did not correctly load community MLX-native checkpoints.

I found three weight-loading problems: convolution weights were transposed twice, decoder keys were not mapped and could leave random weights in place, and the positional-encoding table was initialized to zeros. Quantized checkpoints also needed support for a tokenizer embedded as base64 in the configuration.

The fix detects the checkpoint layout, maps the weights correctly, computes sinusoidal positional encodings, and loads the embedded tokenizer when needed. I added 12 unit tests; the suite had 41 passing tests at the time of the contribution. I also checked full-precision and quantized checkpoints end to end and compared the positional-encoding formula with NeMo. The maintainers reviewed and merged the change.