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