src/transformers/models/whisper/modeling_flax_whisper.py/FlaxWhisperAttention/_merge_heads
class FlaxWhisperAttention:    def _merge_heads(self, hidden_state) -> jnp.ndarray:
        return hidden_state.reshape(hidden_state.shape[:2] + (self.embed_dim,))