[FIXED] Reshape tensor in custom order (PyTorch)

Issue

I have the following tensor

t = torch.tensor([0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14 ,15, 16, 17])

I want to reshape it in the following way:

t_reshape = torch.tensor([[0, 1, 2, 6, 7, 8, 12, 13, 14], 
                          [3, 4, 5, 9, 10, 11, 15, 16, 17]])

Are there ways to efficiently reshape tensors in that fashion?

Solution

You can achieve this by reshaping, transposing and reshaping back:

>>> t.reshape(3,2,-1).transpose(0,1).reshape(2,-1)
tensor([[ 0,  1,  2,  6,  7,  8, 12, 13, 14],
        [ 3,  4,  5,  9, 10, 11, 15, 16, 17]])

Answered By - Ivan