Slice 4d tensor into 4D tensor of smaller subtensors (slice in last 2 dimensions only)

Question

The question is analogous to Slice 2d array into smaller 2d arrays except for the fact that I use tensors (torch) & I have a 4D, not 2D, tensor of the shape eg. (3, 1, 32, 32) - in my case, it is 3 images of size 32x32.

I want to split each tensor of form [i, 0, :, :] into smaller subarrays, so the output would have a shape eg. (3, 16, 8, 8), where each [:, j, :, :] is a small square cut from the original image. I cannot find a way to modify the proposed solution for 4D tensor.

I also tried to just use

subx = x.reshape(3, 16, 8, 8)

but this does not reshape it as I want.

Answer 1

reshape will not work for this purpose. You could look into skimage's view_as_blocks, where the resulting blocks are non-overlapping views of the input array:

from skimage.util.shape import view_as_blocks
view_as_blocks(a, block_shape=(3,1,8,8)).reshape(3, 16, 8, 8)