"Reverse" inverse indices; torch.unique()

messi: <p>I am trying to use <a href="https://pytorch.org/docs/stable/generated/torch.unique.html" rel="noopener nofollow ugc"><code>torch.unique()</code></a>.</p>
<p>Let’s say I have a 2D tensor <code>X = [[1, 2, 3, 4], [1, 2, 3, 4], [3, 4, 5, 6]]</code>. If I use the operation <code>output, inverse_indices = torch.unique(X), sorted=True, return_inverse=True, dim=0)</code>, then <code>output</code> would be  <code> [[1, 2, 3, 4], [3, 4, 5, 6]]</code>  and <code>inverse_indices</code> would be <code>[0, 0, 1]</code>.</p>
<p>That is, <code>inverse_indices</code> is a tensor of the same length as <code>X</code> and tells us the index of elements of <code>X</code> with respect to the tensor <code>output</code>.</p>
<p>Is there a way to obtain the “reverse” of inverse indices? That is, I want to obtain the indices of <code>output</code> with respect to <code>X</code>. For example, in this case, I want to obtain the indices as <code>[0, 2]</code> (or <code>[1, 2]</code>). I understand that this “reverse” inverse index is not unique since multiple elements of <code>X</code> can be mapped to the same element in <code>output</code>.</p>

ptrblck: <p>To get the first index you could use this loop:</p>
<pre><code class="lang-python">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])
</code></pre>

messi: <p>Thanks, this makes sense. But is there any way to vectorize this without loops?</p>

Paul_DUBOIS: <p>It’s a python for loop, so not very efficient, but at least it remains $O(n)$:</p>
<pre><code class="lang-auto">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
</code></pre>

ggoossen: <p>Without using a loop:</p>
<pre><code>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)
indices = torch.scatter_reduce(
    torch.full(size=[output.size(0)], fill_value=inverse_indices.size(0), dtype=torch.int64),
    index=inverse_indices,
    src=torch.arange(inverse_indices.size(0), dtype=torch.int64),
    dim=0,
    reduce="amin",
)
print(f"indices: {indices}")
&gt; indices: tensor([0, 2])
</code></pre>

andrewcaunes: <p><a class="mention" href="/u/ggoossen">@ggoossen</a> function works really well, much faster than others.</p>
<p>By the way it may be minor details at this point but there is a way to make it 10x faster from my testing by not including self :</p>
<pre><code class="lang-auto">def torch_unique_with_indices(tensor):
    """Return the unique elements of a tensor and their indices."""

    unique, inverse_indices = torch.unique(tensor, return_inverse=True)
    indices = torch.scatter_reduce(
        torch.zeros_like(unique, dtype=torch.long, device=tensor.device), 
        dim=0,
        index=inverse_indices,
        src=torch.arange(tensor.size(0), device=tensor.device),
        reduce="amin",
        include_self=False,
    )
    return unique, inverse_indices, indices
</code></pre>
<pre><code class="lang-auto">Time unique
10.1 ms ± 520 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
Time unique with indices include self
156 ms ± 1.61 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)
Time unique with indices not include self
12.3 ms ± 75.7 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
</code></pre>
<p>There is only a 20% overhead compared to torch.unique in my test with a random int array of shape (1000000).<br>
My guess is using torch.zeros_like is faster than torch.full (possible with include_self=False only)</p>