mlx-audio
A merged open-source contribution to mlx-audio, the library for running speech models natively on Apple Silicon. NVIDIA's Canary-1B-v2 recognition model wouldn't load from the community's MLX-native checkpoints; the fix addresses three issues in the loader — a double-transposed convolution, a decoder falling back to random weights, and a zeroed-out positional encoding — and ships with a full test suite, merged upstream.
mlx-audio is a widely-used open-source library for running speech models — text-to-speech and speech recognition — natively on Apple Silicon through Apple’s MLX framework. Its loader for NVIDIA’s Canary-1B-v2 recognition model only understood one specific checkpoint layout (the NeMo-native conversion), so the community-published MLX-native checkpoints on Hugging Face — both full-precision and 8-bit quantized — failed to load entirely.
The loader had three distinct issues with the MLX-native layout. Convolution weights were transposed a second time even though these checkpoints already store them in MLX’s (out, kH, kW, in) layout, which crashed the encoder on a conv2d channel mismatch. The decoder’s weights were dropped because their key names didn’t match the NeMo names the loader expected — so under non-strict loading the decoder fell back to random weights, leaving a model that ran without error but produced incorrect output. And the positional-encoding table was initialized to zeros rather than the sinusoidal values its docstring described, so position information was lost. A separate case: the quantized checkpoint embeds its SentencePiece tokenizer as base64 in config.json instead of shipping a tokenizer.model file, so no tokenizer attached.
The change detects the on-disk layout and routes to the correct key mapping without the extra transpose, computes the sinusoidal positional table to match NeMo, and reads the embedded base64 tokenizer when no model file is present — leaving the original NeMo path unchanged and still covered by its tests. It adds twelve unit tests (full suite: 41 passing) and was verified end-to-end against both the full-precision and quantized checkpoints, with the positional-encoding formula checked against the NeMo reference. Reviewed and merged upstream.