src/transformers/models/t5/modeling_flax_t5.py/FlaxT5Attention/_merge_heads
class FlaxT5Attention:    def _merge_heads(self, hidden_states):
        return hidden_states.reshape(hidden_states.shape[:2] + (self.inner_dim,))