Split a torch tensor using a same-sized tensor of indices

user_9301512 0
Viewed 665

Let's say that I have tensor

t = torch.tensor([1,2,3,4,5])

I want to split it using a same-sized tensor of indices that tells me for each element, in which split it should go.

indices = torch.tensor([0,1,1,0,2])

So that the final result is

splits
[tensor([1,4]), tensor([2,3]), tensor([5])]

Is there a neat way to do this in Pytorch?

EDIT : In general there will be more than 2 or 3 splits.

torch tensor pytorch indexing python
3 Answers

One could do it using argsort for general case:

def mask_split(tensor, indices):
    sorter = torch.argsort(indices)
    _, counts = torch.unique(indices, return_counts=True)
    return torch.split(t[sorter], counts.tolist())


mask_split(t, indices)

Though it might be better to use @flawr answer if this is your real use case (also list comprehension might also be faster as it does not require sorting), something like this:

def mask_split(tensor, indices):
    unique = torch.unique(indices)
    return [tensor[indices == i] for i in unique]
user_108864200

That is indeed possible using logical indexing, you just have to make sure that the index "mask" is made from boolean vales, so in your case

splits = t[indices > 0] , t[indices < 1]

or alternatively you can first cast your tensor indices to have boolean dtype.

user_29131060

On top of other answers, for indexing in pytorch you can directly use index locations to access those elements:

t = torch.tensor([1,2,3,4])
print(t[[0,1,3]])

So you don't really need to store a tensor for indices. You can still store a numpy array with ones and zeros if you want, and then find access indices from that array:

a = np.array([0, 1, 1, 0])
ind_ones = np.argwhere(a == 1).squeeze()
ind_zers = np.argwhere(a == 0).squeeze()
print(t[ind_ones])   # tensor([2, 3])
print(t[ind_zers])   # tensor([1, 4])
user_16071040
Related