Pytorch- Vectorisation of code

I need to perform some operation on a 2D Tensor without using for loops, specifically: Deleting different columns in different rows

2-D Tensor A of shape = [ BS, n*c ]
I need a 2-D Tensor A_out of shape = [ BS, (n-1)*c ]

Essentially, I want to delete c values from each row. These c values will come from different columns for each row.
These specific columns will be decided by the below tensor L:
Shape- [ BS, 1 ] and Values can be integers from ( 0, n-1 )

I want to remove the values between indices c * L[i] and c * ( L[i] + 1 ) for row i

The best I could come up with is as follows. If anyone can help with removing the for loop in this, it’ll be really appreciated:

A_out = torch.zeros( BS, (n-1) * c)

for i in range( BS ):
A_out[i] = torch.cat( [ A[ i , :c * L[i] ], A[ i , c * ( L[i]+1 ):] ] )

(‘gather’ seemed like the appropriate option at first, but creating the 2-D tensor of indexes to gather using L also seemed equally costly i.e. for loop needed)

Hi Ayush,

Welcome to the forums. And thanks for the very enjoyable morning puzzle!

Here’s one way to do it:

BS = 6
c = 2
n = 5
A = torch.arange(BS*c*n).view(BS,c*n)
A
tensor([[ 0,  1,  2,  3,  4,  5,  6,  7,  8,  9],
        [10, 11, 12, 13, 14, 15, 16, 17, 18, 19],
        [20, 21, 22, 23, 24, 25, 26, 27, 28, 29],
        [30, 31, 32, 33, 34, 35, 36, 37, 38, 39],
        [40, 41, 42, 43, 44, 45, 46, 47, 48, 49],
        [50, 51, 52, 53, 54, 55, 56, 57, 58, 59]])

L = torch.LongTensor([1,2,3,0,4,2]) #Columns to delete

In PyTorch, it’s generally simpler to work with whole columns rather than sequences inside a row. And, generally, a boolean mask is a powerful way to delete elements and slices.

A = A.view(BS,n,c)  # Areas to be deleted become whole column slices instead of BS sequences of length c.
mask = torch.ones((BS,n), dtype=torch.bool)  #We will delete whole columns using a boolean mask

# Mark the slices to be deleted.
mask[range(BS),L] = False
# (There may be a more efficient way to do this than using range, 
# hidden somewhere in the arcana of broadcasting and indexing.)

mask
tensor([[ True, False,  True,  True,  True],
        [ True,  True, False,  True,  True],
        [ True,  True,  True, False,  True],
        [False,  True,  True,  True,  True],
        [ True,  True,  True,  True, False],
        [ True,  True, False,  True,  True]])

#Delete the column slices and reshape back to 2 dimensions.
A[mask].view(BS,-1) #Here, mask is broadcast over the c dimension.

tensor([[ 0,  1,  4,  5,  6,  7,  8,  9],
        [10, 11, 12, 13, 16, 17, 18, 19],
        [20, 21, 22, 23, 24, 25, 28, 29],
        [32, 33, 34, 35, 36, 37, 38, 39],
        [40, 41, 42, 43, 44, 45, 46, 47],
        [50, 51, 52, 53, 56, 57, 58, 59]])

HTH, Malcolm

Hi Ayush,

Did this approach work for you?

Hi Malcolm.
Thanks for the prompt response and also for checking back.
On just reading, it seems to be exactly what I was looking for.
Because of some personal constraints, I haven’t been able to check the correctness and timing improvement of the solution yet. I’ll get back with those when possible for me.
Thanks a lot :slight_smile: