src/transformers/models/ctrl/modeling_tf_ctrl.py/TFMultiHeadAttention/split_into_heads
class TFMultiHeadAttention:    def split_into_heads(self, x, batch_size):
        x = tf.reshape(x, (batch_size, -1, self.num_heads, self.depth))
        return tf.transpose(x, perm=[0, 2, 1, 3])