src/transformers/models/blenderbot_small/modeling_flax_blenderbot_small.py/FlaxBlenderbotSmallAttention/_merge_heads
class FlaxBlenderbotSmallAttention:    def _merge_heads(self, hidden_states):
        return hidden_states.reshape(hidden_states.shape[:2] + (self.embed_dim,))