src/transformers/models/ctrl/modeling_ctrl.py/MultiHeadAttention/split_into_heads
class MultiHeadAttention:    def split_into_heads(self, x, batch_size):
        x = x.reshape(batch_size, -1, self.num_heads, self.depth)
        return x.permute([0, 2, 1, 3])