src/transformers/models/big_bird/modeling_flax_big_bird.py/FlaxBigBirdSelfAttention/_merge_heads
class FlaxBigBirdSelfAttention:    def _merge_heads(self, hidden_states):
        return hidden_states.reshape(hidden_states.shape[:2] + (self.config.hidden_size,))