Create a Image stack in PyTorch
I have a Bx3xHxW image tensor in PyTorch and wish to create a Bx3FxHxW image stack of this image where F=64.
The image stack is formed by shifting the original image right. That is if the original image is to be shifted by 2 pixels to the right, the two left-most columns in new image will be 0 and the third-last column of the original image will become the last column of the new shifted image.
The image stack is formed by shifting right the original image f times where f varies from 0 to F-1.
How to achieve this in PyTorch in the most effecient manner using no or minimal number of for loops? A numpy code will also work as the two languages are quite compatible.

Torch (as well as numpy) provides torch.roll function, by padding with zeros first, rolling and then slicing the result you can achieve your right shift.
Here's a numpy version:
import numpy as np

X = np.random.rand(4,3,28,28)
Z = np.zeros((4,3,28,28))
XZ = np.concatenate([X,Z],axis=-1)

res = []
shift = 2
F = 28//shift
for i in range(0,F,shift):
    res.append(np.roll(XZ,i,-1)[:,:,:,:28])

res = np.concatenate(res,axis=1)

Here's a 1D example for better understanding:
x = np.arange(10)
z = np.zeros(10)
xz = np.concatenate([x,z],axis=-1)

for i in range(0,10,2):
    print(np.roll(xz,i,-1)[:10])

[0. 1. 2. 3. 4. 5. 6. 7. 8. 9.]
[0. 0. 0. 1. 2. 3. 4. 5. 6. 7.]
[0. 0. 0. 0. 0. 1. 2. 3. 4. 5.]
[0. 0. 0. 0. 0. 0. 0. 1. 2. 3.]
[0. 0. 0. 0. 0. 0. 0. 0. 0. 1.]

I hope this is what you are looking for.
