How to use a 1D tensor as an index to slice a 2D tensor?

1

Normally, you can slice a 2D tensor like this slice = t[:, :k] where k is an integer. Is it possible to do something like this but with k being a 1-dimensional vector of integers with the number of items that I want to obtain for each row?

Masking the items with 0’s or NaN would also be fine.

For example:

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

# perform some operations and the result should be
# 1 - -
# 2 2 2
# 3 3 -
2

I could think of the following way:

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

h, w = t.shape

# create a mask of ones with size h, w+1
mask = torch.ones(h, w+1)

# set other elements to 0, depending on length of each row
mask[torch.arange(h), k.long()] = 0.
mask = mask.cumprod(dim=1)

# multiply the mask
result = t * mask[:, :-1]
print(result)

 tensor([[1., 0., 0.],
        [2., 2., 2.],
        [3., 3., 0.]])
3

I ended up doing something similar. But, instead of multiplying, I just did t[mask] = 0