"Reverse" inverse indices; torch.unique()

#1

I am trying to use torch.unique().

Let’s say I have a 2D tensor X = [[1, 2, 3, 4], [1, 2, 3, 4], [3, 4, 5, 6]]. If I use the operation output, inverse_indices = torch.unique(X), sorted=True, return_inverse=True, dim=0), then output would be [[1, 2, 3, 4], [3, 4, 5, 6]] and inverse_indices would be [0, 0, 1].

That is, inverse_indices is a tensor of the same length as X and tells us the index of elements of X with respect to the tensor output.

Is there a way to obtain the “reverse” of inverse indices? That is, I want to obtain the indices of output with respect to X. For example, in this case, I want to obtain the indices as [0, 2] (or [1, 2]). I understand that this “reverse” inverse index is not unique since multiple elements of X can be mapped to the same element in output.

#2

To get the first index you could use this loop:

X = torch.tensor([[1, 2, 3, 4], [1, 2, 3, 4], [3, 4, 5, 6]])
output, inverse_indices = torch.unique(X, sorted=True, return_inverse=True, dim=0)

for o in output:
    print((o == X).all(1).nonzero()[0])
1 Like
#3

Thanks, this makes sense. But is there any way to vectorize this without loops?

#4

It’s a python for loop, so not very efficient, but at least it remains $O(n)$:

output, inverse_indices = torch.unique(X, return_inverse=True)

indices = torch.zeros_like(output, dtype=torch.int)
for i in range(indices.shape[0]):
    indices[inverse_indices[i]] = i