3D tensor transformation


Imagine that I have a cube of ones:

cube = torch.Tensor(

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(

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]):
    return doc

then just use torch.sum along dim0…

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