Implement beam search decoding for Whisper

[muaz978]

What this adds

DecodingOptions already accepts beam_size and patience, and most of the machinery for group decoding exists (batch tiling for n_group > 1, MaximumLikelihoodRanker, Inference.rearrange_kv_cache), but requesting a beam size raises NotImplementedError. This PR implements the missing BeamSearchDecoder, ported from the reference implementation in openai/whisper and adapted to MLX.

Implementation notes

• Device-side top-k: candidate selection runs on the GPU with mx.argpartition + mx.take_along_axis, so only (beam_size + 1) token ids and log probabilities per beam cross to the host each step, instead of a (n_batch, n_vocab) matrix. The beam bookkeeping itself stays in python, like the reference implementation.
• Host-side prefix reuse: candidate prefixes are kept between steps instead of pulling the whole growing token matrix off the device every step.
• Rectangular finalize: finished sequences are padded with EOT so the result matches what DecodingTask.run expects before it truncates each row at the first EOT.
• Cross-attention safe cache rearrange: rearrange_kv_cache reindexed every cached array, but the cache mixes self-attention entries (batch = n_audio * n_group) with cross-attention entries (batch = n_audio, broadcast over the group). Indexing the cross-attention arrays with beam indices reads out of bounds, which MLX does not bounds-check on the GPU, silently corrupting the decode into nondeterministic garbage. The rearrange now only touches arrays whose batch dimension matches the selection; the cross-attention rows are identical across beams anyway. (This also future-proofs best_of style uses that might want to rearrange.)

Testing

On Apple Silicon (M series), with Arabic and English speech:

• beam_size=5 at temperature=0 is deterministic across runs and returns correct transcriptions (previously: NotImplementedError).
• Works together with word_timestamps=True, patience, and the temperature fallback chain (beam_size applies at t == 0, best_of sampling above, matching openai/whisper).
• decode still rejects beam_size + best_of given together, and patience without beam_size.
• Wall clock with whisper-large-v3-turbo on a 44 s clip: beam_size=5 was within a few percent of greedy decoding. Smaller models pay relatively more (about 2x on whisper-small) because the per-step bookkeeping is a larger share of their forward pass.
• Formatted with black.

Beam search noticeably improves difficult, accented, and dialectal speech compared to greedy decoding (it is what faster-whisper and openai/whisper use by default with beam_size=5), so this closes one of the last accuracy gaps between mlx_whisper and those implementations.