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