EDITED: added channels (making things even harder )
Hey there,
I have a tensor of shape [batch, C, H, W] which is a batch of images with 3 channels
batch = torch.rand((2,3,64,64))

I’m manipulating this batch, and my goal is to return to this shape after the manipulation was done.
The manipulation is to slice each image to 16x16 squares and represent each square as a flattened vector, and lastly, since each image was sliced, we’ll take all the slices and represent them as a sequence.
The manipulation is as follow:
patched_squares = batch.unfold(2,16,16).unfold(3,16,16)

which should slice each image-channel to 16x16 squares, so the output shape will be:

torch.Size([2, 3, 4, 4, 16, 16])

meaning, a batch of size 2, 3 channels, with 4 rows and 4 columns that consist of 16x16 squares
next I’m flattening the 16x16 squares to a vector, like so:
patched_squares = patched_squares.flatten(start_dim=4, end_dim=5)

now the size is:

torch.Size([2, 3, 4, 4, 256])

meaning, a batch of size 2, 3 channels, with 4 rows and 4 columns that consist of 256 flattened vector
Next, I’m flattening the 4x4 to be a sequence:
patched_squares = patched_squares.flatten(start_dim=2, end_dim=3)

which gives a batch of size 2, 3 channels, with sequence length = 16, where each sample in the sequence is a 256 vector

torch.Size([2, 3, 16, 256])

Lastly, I’m concatenating the channels of each “square” like so:
patched_squares = torch.stack([torch.cat([p[0],p[1],p[2]], dim=1) for p in patched_squares])


torch.Size([2, 16, 768])

The full manipulation code is:
batch = torch.rand((2,3,64,64))
patched_squares = batch.unfold(2,16,16).unfold(3,16,16).flatten(start_dim=4, end_dim=5).flatten(start_dim=2, end_dim=3)
patched_squares = torch.stack([torch.cat([p[0],p[1],p[2]], dim=1) for p in patched_squares])

And now to the challenge itself 
After running some computation on patched_squares, I would like to “refold” it back to how it was,
so from [2, 16, 768] back to [2, 3, 64, 64], while maintaining the actual original positions.
I’m having a hard time to refold, I would appreciate any help in this manner.
Roy