Slicing 2D tensor using list of start/end indices

1

Given a 2D tensor x and 2D tensors of start and end indices, how can I get a tensor with the slices of x? I am looking for a vectorized function that can do this:

>>> x = torch.arange(0, 20).view(4, 5)
tensor([[ 0,  1,  2,  3,  4],
        [ 5,  6,  7,  8,  9],
        [10, 11, 12, 13, 14],
        [15, 16, 17, 18, 19]])
>>> starts = torch.tensor([[0,0], [0,3], [2,1]])
tensor([[0,0], [0,3], [2,1]])
>>> ends = torch.tensor([[0,2], [0,5], [2,5]])
tensor([[0,2], [0,5], [2,5]])
>>> slices = some_function(x, starts, ends)
tensor([[0, 1], [3, 4],
           [],
           [11, 12, 13, 14],
           []])

starts and ends have the same length and each slice given by the pair start[i], end[i] only runs over the same zeroth dimension. Any help would be appreciated :slight_smile:

2

First, your output is not a tensor because each row don’t have the same lenght, if they have the same length, you can refer here, otherwise I’m afraid it has to be a loop to solve it