reshape tensor in a determined way

user_15577667 0
Viewed 38

I have some tensor x3. I got it in the following way:

x = torch.tensor([0, 0, 0, 0, 1, 0, 0, 0, 0])
x2 = torch.stack(5 * [x], 0)
x2 = x2.reshape(-1)
x3 = torch.stack(4 * [x2], 0)
x3 = torch.stack(6 * [x3], -1)
x3 = torch.stack(7 * [x3], -1)

In short it means that

x[0, :9, 0, 0] = [0, 0, 0, 0, 1, 0, 0, 0, 0]
x[0, 9:18, 0, 0] = [0, 0, 0, 0, 1, 0, 0, 0, 0]

and so on.
Then I want to reshape it that every nine values of the 1st dimension go to the new dimension. In other words, I want x3[0, 0, 0, 0, :] to give me tensor([0, 0, 0, 0, 1, 0, 0, 0, 0])

I tried to do:

x3.reshape(4, 5, 6, 7, 9)[0, 0, 0, 0, :]
tensor([0, 0, 0, 0, 0, 0, 0, 0, 0])
x3.reshape(4, 9, 6, 7, 5).transpose(1, -1)[0, 0, 0, 0, :]
tensor([0, 0, 0, 0, 0, 0, 0, 0, 1])

As you see, none of it gives me the right answer

UPD: added x3 = torch.stack(7 * [x3], -1)

pytorch python
1 Answers

If you want to modify the 1st dimension and create a new dimension at the end, you need to first move that dimension to last by using permute. Something like this should do the trick:

xpermuted = x3.permute(0, 2, 3, 1)
xreshaped = xpermuted.reshape(xpermuted.shape[0], xpermuted.shape[1], xpermuted.shape[2], int(xpermuted.shape[3] / 9), 9)

print(xreshaped[0, 0, 0, 0, :])     # tensor([0, 0, 0, 0, 1, 0, 0, 0, 0])
print(xreshaped[0, 0, 0, 1, :])     # tensor([0, 0, 0, 0, 1, 0, 0, 0, 0])
print(xreshaped[0, 0, 0, 2, :])     # tensor([0, 0, 0, 0, 1, 0, 0, 0, 0])

After that, you can restore the initial dimension order by using permute again if you need the original order of dimensions:

xrestored = xreshaped.permute(0, 3, 1, 2, 4)
print(xrestored.shape)              # torch.Size([4, 5, 6, 7, 9])

Technically you don't have to move the first dimension to last initially, you can also do the reverse first reshape, then permute. Actually, now that I think about it, this is better since it has one less permute:

xreshaped = x3.reshape(x3.shape[0], int(x3.shape[1]/9), 9, x3.shape[2], x3.shape[3])
xrestored = xreshaped.permute(0, 1, 3, 4, 2)

print(xrestored.shape)              # torch.Size([4, 5, 6, 7, 9])
print(xrestored[0, 0, 0, 0, :])     # tensor([0, 0, 0, 0, 1, 0, 0, 0, 0])
print(xrestored[0, 1, 0, 0, :])     # tensor([0, 0, 0, 0, 1, 0, 0, 0, 0])
user_16071040
Related