"Reverse" map of `inverse` in torch.unique?

yatorho: <p>Hi all.</p>
<p>I want to implement a function like “reverse” map of <code>inverse</code> value of <code>torch.unique</code>.</p>
<p>For example, <code>torch.unique</code> function can return the unique values <code>uniques</code> of long type input <code>x</code> and a inverse tensor  that is a mapping from <code>uniques</code> to <code>x</code>.</p>
<pre><code class="lang-auto">x = torch.LongTensor([9, 10, 9, 9, 10, 9])
uniques, inverse, counts = torch.unique(x, return_inverse=True, return_counts=True)
# uniques = [9, 10]
# inverse = [0, 1, 0, 0, 1, 0] 
# counts = [4, 2]
print((uniques[inverse] == x).all())  # True
</code></pre>
<p>For my question, is there some efficient way to get “reverse” inverse <code>back_map</code> that mapping from <code>x</code> to <code>uniques</code>?</p>
<pre><code class="lang-auto">def reverse_unique(x): ...

uniques, inverse, counts, back_map = reverse_unique(x)
# uniques = [9, 10]
# inverse = [0, 1, 0, 0, 1, 0] 
# counts = [4, 2]
# back_map = [0, 2, 3, 5, 1, 4]
print((x[back_map] == uniques.repeat_interleave(counts)).all()) # True
</code></pre>
<p>In above code, the <code>back_map</code> maps the values of  <code>inverse</code> to the position of the input <code>x</code>.</p>
<p>I know it not a difficult thing implements this function with python loop, but in my case where the input <code>x</code> size can reach e8, so the time overhead is intolerable.</p>
<p>So, Is there any high level implementation using pytorch api or the cuda kernel(I tried to use cuda extension to parallelize it but my cuad kernel is slow extremely <img src="https://discuss.pytorch.org/images/emoji/apple/sob.png?v=12" title=":sob:" class="emoji" alt=":sob:" loading="lazy" width="20" height="20">)?</p>
<pre><code class="lang-auto">__global__ void unique_back_map_kernel(
    int32_t num_uni,
    int32_t num_x,
    int64_t* __restrict__ uniques,
    int64_t* __restrict__ cumsum_counts,
    int64_t* __restrict__ x,
    int64_t* __restrict__ out) {
  int32_t n = blockIdx.x * blockDim.x + threadIdx.x;
  if (n &gt;= num_uni) {
    return;
  }

  size_t counts = 0;
  auto idx = __ldg(&amp;cumsum_counts[n]);

#pragma unroll
  for (int64_t i = 0; i &lt; num_x; ++i) {
    if (x[i] == uniques[n]) {
      out[idx + counts] = i;
      counts++;
    }
  }
}

std::tuple&lt;at::Tensor, at::Tensor, at::Tensor, at::Tensor&gt; reverse_unique(at::Tensor x) {
  TORCH_CHECK(x.dtype() == at::kLong &amp;&amp; x.device().is_cuda());
  auto [uniques, inverse, counts] = at::_unique2(x, false, true, true);
  counts.cumsum_(0);
  // python: cumsum_counts = torch.cat([torch.tensor([0]), cumsum_counts[:-1]])
  auto cumsum_counts = at::cat({at::zeros({1}, counts.options()), counts.slice(0, 0, -1)});

  auto back_map = at::empty_like(x);

  int32_t threads = (uniques.numel() &gt; 256) ? 256 : 32;
  int32_t blocks = (uniques.numel() + threads - 1) / threads;
  unique_back_map_kernel&lt;&lt;&lt;blocks, threads, 0, c10::cuda::getCurrentCUDAStream()&gt;&gt;&gt;(
      uniques.numel(),
      x.numel(),
      uniques.data_ptr&lt;int64_t&gt;(),
      cumsum_counts.data_ptr&lt;int64_t&gt;(),
      x.data_ptr&lt;int64_t&gt;(),
      back_map.data_ptr&lt;int64_t&gt;());

  return std::make_tuple(uniques, inverse, cumsum_counts, back_map);
}
</code></pre>

yatorho: <p>The problem has been solved.</p>
<p><code>torch.sort</code> not only returns the sorted values but also the permutation of indices for getting back the original order. This permutation is exactly <code>back_map</code>.</p>
<pre><code class="lang-auto">_, back_map = torch.sort(x)
</code></pre>
<p>So, the <code>reverse_unique</code> function can be implemented as:</p>
<pre><code class="lang-auto">def reverse_unique(x):
    x_sorted, back_map = torch.sort(x)
    uniques, inverse, counts = torch.unique_consecutive(
        x_sorted, return_inverse=True, return_counts=True
    )

    return uniques, inverse, counts, back_map
</code></pre>