fix: Gemma 4 audio — mel preprocessing, weight loading, feature extractor

[stephencox-ict]

Overview

Fix four bugs preventing Gemma 4 audio (E2B/E4B) from producing correct transcriptions. These were discovered while implementing and debugging the Gemma 4 audio conformer encoder for llama.cpp (ggml-org/llama.cpp#21421), where we traced tensors between PyTorch and C++ to identify divergences.

Fixes

1. Missing semicausal left-padding (audio_feature_extractor.py)

The HF reference prepends frame_length // 2 (160) zero samples before the unfold, centering the first frame at t=0. Without this:

• The mel spectrogram is misaligned vs what the model was trained on
• The frame count is systematically fewer, causing the token count mismatch that triggers the broadcast shapes error (Gemma4 Value Error - Broadcast Shapes #923)

2. Wrong Hann window formula (audio_feature_extractor.py)

Used cos(2π(n+0.5)/N) instead of the correct periodic Hann cos(2πn/N). The +0.5 phase shift produces meaningfully different spectral values. Verified against HF Transformers reference which uses window_function(frame_length) = periodic Hann.

3. sanitize() double-nests model. prefix (gemma4.py) — fixes #912

HF checkpoint keys like model.language_model.model.embed_tokens.weight become language_model.model.embed_tokens.weight after stripping the outer model. prefix — this already matches the MLX attribute path. The unconditional .model. insertion created language_model.model.model.*, so all language model weights loaded as zero.

4. Feature extractor not instantiated (processing_gemma4.py) — fixes #903

The feature extractor was only created when processor_config.json contains a "feature_extractor" key, which standard HF checkpoints don't include. Now instantiates with USM defaults unconditionally (the parameters are fixed for all Gemma 4 audio models: 16kHz, 128 mel bins, 20ms frame, 10ms hop, HTK scale).

Additional information

Fixes #903, #912, #923

Note: I don't have an Apple Silicon machine to test this on, so these fixes are based on code review and our experience fixing the same issues in the llama.cpp C++ implementation. The mel preprocessing fixes (semicausal padding, Hann window) were verified against PyTorch with mel cosine similarity of 0.9998.

Requirements

• I have read and agree with the contributing guidelines
• AI usage disclosure: YES — Claude Code was used to assist with code review and identifying the bugs. All fixes are based on manual tracing against the HuggingFace Transformers reference implementation.

[Blaizzy]

Thanks!