src/transformers/models/wav2vec2/modeling_flax_wav2vec2.py/FlaxWav2Vec2Attention/_split_heads
class FlaxWav2Vec2Attention:    def _split_heads(self, hidden_states):
        return hidden_states.reshape(hidden_states.shape[:2] + (self.num_heads, self.head_dim))