generate 1D tensor as unique index of rows of an 2D tensor

user_5319143 0
Viewed 83

Let's say we transform a 2D tensor to a 1D tensor by giving each, different row a different index, from 0 to the number of rows - 1.

[[1,2],[1,3],[1,4]] -> [0,1,2]

But if there are same rows, then we repeate the index, like this below.

[[1,2],[1,2],[1,4]] -> [0,0,2]

[[1,2],[1,3],[1,2]] -> [0,1,0]

How to implement this on PyTorch?

tensor pytorch python
1 Answers

You can do so using torch.Tensor.unique and returning the inverse which provides the indices out of the box:

>>> _, i = x.unique(dim=0, return_inverse=True)

>>> i # first example
tensor([0, 1, 2])
user_63313690
Related