src/transformers/models/xlm_roberta/modeling_flax_xlm_roberta.py/FlaxXLMRobertaSelfAttention/_split_heads
class FlaxXLMRobertaSelfAttention:    def _split_heads(self, hidden_states):
        return hidden_states.reshape(hidden_states.shape[:2] + (self.config.num_attention_heads, self.head_dim))