Decoder state one decode owns (self-attention K/V and the cross K/V from the encoder output).
The graph that made it is the only one that reads it; the transcription loop only advances,
reorders and releases it.
Beam reorder: keep rows indices (mx.take along axis 0). Cross K/V
rows are identical across beams and are left untouched when they were
broadcast from one audio (shape[0] === 1).
Decoder state one decode owns (self-attention K/V and the cross K/V from the encoder output). The graph that made it is the only one that reads it; the transcription loop only advances, reorders and releases it.