src/transformers/models/cvt/modeling_tf_cvt.py/TFCvtSelfAttention/rearrange_for_multi_head_attention
class TFCvtSelfAttention:    def rearrange_for_multi_head_attention(self, hidden_state: tf.Tensor) -> tf.Tensor:
        batch_size, hidden_size, _ = shape_list(hidden_state)
        head_dim = self.embed_dim // self.num_heads
        hidden_state = tf.reshape(hidden_state, shape=(batch_size, hidden_size, self.num_heads, head_dim))
        hidden_state = tf.transpose(hidden_state, perm=(0, 2, 1, 3))
        return hidden_state