src/transformers/models/longt5/modeling_flax_longt5.py/FlaxLongT5LocalAttention/_split_heads
class FlaxLongT5LocalAttention:    def _split_heads(self, hidden_states):
        return hidden_states.reshape(hidden_states.shape[:2] + (self.n_heads, self.key_value_proj_dim))