Torch.split with multiple start and end indices

1

Given a tensor [0,1,2,3,4,5,6,7,8,9], and a slices tensor of shape [N_SLICES, 2], e.g. [[0,2], [4,7]]. I would like to gather all the respective slices without using for loop (for every [start,end] pair and appending the slices). Is there a method that achieves this?

In the example above, the result should be [0,1,4,5,6]. I believe torch.split only splits into chunks given that the resulting split tensor has the same shape as the original tensor. In my case, the resulting split tensor have less elements than the original tensor.

1 Like
2

Hi Nielsen!

Yes, consider:

>>> import torch
>>> torch.__version__
'2.7.0+cu128'
>>> t = torch.tensor ([0, 1, 2, 3, 4, 5, 6, 7, 8, 9])
>>> s = torch.tensor ([[0, 2], [4, 7]])
>>> ind = torch.arange (len (t)).expand (s.size (0), -1)
>>> mask = torch.logical_and (ind >= s[:, :1], ind < s[:, 1:]).any (0)
>>> t[mask]
tensor([0, 1, 4, 5, 6])

Best.

K. Frank

1 Like