undo torch.chunk to rebuild image in right order
I am trying to process a 3D image in chunks (non-overlapping windows). Once this is done I want to put the chunks back together in the right order.
I have been chunking the image as below:
tens = torch.tensor(range(64))
tens = tens.view((4,4,4))
print(tens)
&gt;&gt;&gt;tensor([[[ 0,  1,  2,  3],
         [ 4,  5,  6,  7],
         [ 8,  9, 10, 11],
         [12, 13, 14, 15]],

        [[16, 17, 18, 19],
         [20, 21, 22, 23],
         [24, 25, 26, 27],
         [28, 29, 30, 31]],

        [[32, 33, 34, 35],
         [36, 37, 38, 39],
         [40, 41, 42, 43],
         [44, 45, 46, 47]],

        [[48, 49, 50, 51],
         [52, 53, 54, 55],
         [56, 57, 58, 59],
         [60, 61, 62, 63]]])

tens = torch.chunk(tens,2, -1)
tens = torch.stack(tens)
tens = torch.chunk(tens,2, -2)
tens = torch.concat(tens)
tens = torch.chunk(tens,2, -3)
tens = torch.concat(tens)
print(torch.shape)
&gt;&gt;&gt;torch.Size([8, 2, 2, 2])

Then I want to put it back together in the original order
tens = tens.view([4,4,2,2])
tens = tens.view([2,4,4,2])
tens = tens.view([4,4,4])
print(tens)
&gt;&gt;&gt;tensor([[[ 0,  1,  4,  5],
         [16, 17, 20, 21],
         [ 2,  3,  6,  7],
         [18, 19, 22, 23]],

        [[ 8,  9, 12, 13],
         [24, 25, 28, 29],
         [10, 11, 14, 15],
         [26, 27, 30, 31]],

        [[32, 33, 36, 37],
         [48, 49, 52, 53],
         [34, 35, 38, 39],
         [50, 51, 54, 55]],

        [[40, 41, 44, 45],
         [56, 57, 60, 61],
         [42, 43, 46, 47],
         [58, 59, 62, 63]]])

and I can't figure out how to get the elements in the right order. I realise I probably missed something in the docs or something else obvious but I can't find it. Any ideas?

Operator torch.chunk doesn't reduce dimensions so its inverse is torch.cat, not torch.stack.
Here are the transforms with the corresponding inverse operations:

Splitting the last dimension in two:

&gt;&gt;&gt; chunks = tens.chunk(chunks=2, dim=-1)
&gt;&gt;&gt; torch.cat(chunks, dim=-1)


Splitting the second dimension into two:

&gt;&gt;&gt; chunks = tens.chunk(chunks=2, dim=-2)
&gt;&gt;&gt; torch.cat(chunks, dim=-2)


Splitting the first dimension into two:

&gt;&gt;&gt; chunks = tens.chunk(chunks=2, dim=-3)
&gt;&gt;&gt; torch.cat(chunks, dim=-3)




If you want to invert the whole sequence, you just have to keep in mind that torch.cat is the reverse operation of torch.chunk:
&gt;&gt;&gt; tens = torch.cat(tens.chunk(2), dim=-3)
&gt;&gt;&gt; tens = torch.cat(tens.chunk(2), dim=-2)
&gt;&gt;&gt; tens = torch.cat(tens.chunk(2), dim=-1)

