Vectorizing assignment of a tensor to a slice in PyTorch
I'm trying to vectorize a slice assignment of the form

for i in range(a.shape[1]):
    for j in range(a.shape[2]):
        a[:,i,j,:,i:i+b.shape[2],j:j+b.shape[3]] = b


where b itself is an array. This is because the nested Python loop is too inefficient and is taking up most of the runtime. Is there a way to do this?

For a simpler case, consider the following:

for i in range(a.shape[1]):
    a[:,i,:,i:i+b.shape[2]] = b


This is what b and a might look like:



You can see the diagonal, "sliding" structure of the resulting matrix.

We can leverage np.lib.stride_tricks.as_strided based scikit-image's view_as_windows to get sliding windowed views into a 0s padded version of the input and being a view would be efficient on memory and performance.  More info on use of as_strided based view_as_windows.

Hence, for the simpler case, it would be -

from skimage.util.shape import view_as_windows

def sliding_2D_windows(b, outshp_axis1):
    # outshp_axis1 is desired output's shape along axis=1
    n = outshp_axis1-1
    b1 = np.pad(b,((0,0),(0,0),(n,n)),'constant')
    w_shp = (1,b1.shape[1],b.shape[2]+n)
    return view_as_windows(b1,w_shp)[...,0,::-1,0,:,:]


Sample run -

In [192]: b
Out[192]: 
array([[[54, 57, 74, 77],
        [77, 19, 93, 31],
        [46, 97, 80, 98]],

       [[98, 22, 68, 75],
        [49, 97, 56, 98],
        [91, 47, 35, 87]]])

In [193]: sliding_2D_windows(b, outshp_axis1=3)
Out[193]: 
array([[[[54, 57, 74, 77,  0,  0],
         [77, 19, 93, 31,  0,  0],
         [46, 97, 80, 98,  0,  0]],

        [[ 0, 54, 57, 74, 77,  0],
         [ 0, 77, 19, 93, 31,  0],
         [ 0, 46, 97, 80, 98,  0]],

        [[ 0,  0, 54, 57, 74, 77],
         [ 0,  0, 77, 19, 93, 31],
         [ 0,  0, 46, 97, 80, 98]]],


       [[[98, 22, 68, 75,  0,  0],
         [49, 97, 56, 98,  0,  0],
         [91, 47, 35, 87,  0,  0]],

        ....

        [[ 0,  0, 98, 22, 68, 75],
         [ 0,  0, 49, 97, 56, 98],
         [ 0,  0, 91, 47, 35, 87]]]])

