3D tensor transformation

Hello

Imagine that I have a cube of ones:

cube = torch.Tensor(
[[[1,1,1],[1,1,1],[1,1,1]],
[[1,1,1],[1,1,1],[1,1,1]]
[[1,1,1],[1,1,1],[1,1,1]]])

I want to first “shift” each “layer”/2D Array contained in the cube along the dim1 axis, to obtain something like that (I then have to fill the “emptiness” by zeros):

cube = torch.Tensor(
[[[1,1,1,0,0],[1,1,1,0,0],[1,1,1,0,0]],
[[0,1,1,1,0],[0,1,1,1,0],[0,1,1,1,0]]
[[0,0,1,1,1],[0,0,1,1,1],[0,0,1,1,1]]])

Finally add the three layers of the cube to obtain a (3*5) array:

cube = torch.Tensor([[1,2,3,2,1],[1,2,3,2,1],[1,2,3,2,1]])

I can make a drawing to better explain the shifting part if I’m not clear.

I want to do it with any (x,y,z) tensor in the easier way. Can someone help me ?

If someone intersted I resolved this by doing like that:

def threeto2ddoc(doc):
    doc = torch.cat((doc, torch.zeros(doc.shape[0],doc.shape[1],doc.shape[0]-1)),dim=2)
    for i in range(doc.shape[0]):
        for j in range(doc.shape[1]):
            doc[i,j,:]=torch.cat((doc[i,j,-i:],doc[i,j,0:-i]),dim=0)
    return doc

then just use torch.sum along dim0…

how about a PR if the solution is still not realized :wink: