How to unfold tensor properly to split a large tensor into blocks?

Hi all, so this is a follow-up to a very similar question I had recently, except specifically relating to the unfolding function, since I can’t seem to figure out how to use it properly to split a large tensor into smaller blocks (as if one were running a 3D filter across the tensor). Hopefully this drawing below better clarifies what I’m trying to achieve:


First off, I’m generating a 4x4x4x3 tensor of 3D coordinate points using the following code:

shape = (4,)*3
size = shape[0] * shape[1] * shape[2]

pxs = torch.linspace(-1, 1, shape[0])
pys = torch.linspace(-1, 1, shape[1])
pzs = torch.linspace(-1, 1, shape[2])

pxs = pxs.view(-1, 1, 1).expand(*shape).contiguous().view(size)
pys = pys.view(1, -1, 1).expand(*shape).contiguous().view(size)
pzs = pzs.view(1, 1, -1).expand(*shape).contiguous().view(size)
points = torch.stack([pxs, pys, pzs], dim=1)

points = torch.reshape(points, (eval_dim, eval_dim, eval_dim, 3))

Next, I’m using the following unfolding operations to split this tensor into eight 2x2x2 blocks:

patches = rotated_grids_3d.unfold(2, 2, 2).unfold(1, 2, 2).unfold(0, 2, 2)
patches = patches.contiguous().view(-1, 2, 2, 2, 3)

However, no matter the different unfolding orderings I try, I can’t seem to get the blocks to be nearby points. For example, 2x2x2 block should be the back, top-left corner of the cube, with coordinate points as follows:

print(rotated_grids_4_3d[0, 0, 0, :])
print(rotated_grids_4_3d[0, 0, 1, :])
print(rotated_grids_4_3d[0, 1, 0, :])
print(rotated_grids_4_3d[0, 1, 1, :])
print(rotated_grids_4_3d[1, 0, 0, :])
print(rotated_grids_4_3d[1, 0, 1, :])
print(rotated_grids_4_3d[1, 1, 0, :])
print(rotated_grids_4_3d[1, 1, 1, :])

Then the next block would be two points moved over, etc. etc. However, my unfolding results seem to give arbitrarily located points that don’t themselves form cubes, and there’s sometimes repeating points, and I haven’t been able to figure out what I’m doing wrong. Any help would be appreciated!