我想实现一个类似于torch.unique的inverse值的“反向”Map的函数。
例如，torch.unique函数可以返回长类型输入x的唯一值uniques和从uniquesMap到x的逆Tensor。

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

字符串
对于我的问题，是否有一些有效的方法来获得从xMap到uniques的“反向”逆back_map？

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

型
在上面的代码中，back_map将inverse的值Map到输入x的位置。
我知道用python循环实现这个函数并不困难，但在我的例子中，输入的x大小可以达到e8，所以时间开销是无法忍受的。
那么，有没有使用pytorch API或cuda内核的高级实现（我试图使用cuda扩展来并行化它，但我的cuad内核非常慢：sob：）？

__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 >= num_uni) {
    return;
  }

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

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

std::tuple<at::Tensor, at::Tensor, at::Tensor, at::Tensor> reverse_unique(at::Tensor x) {
  TORCH_CHECK(x.dtype() == at::kLong && 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() > 256) ? 256 : 32;
  int32_t blocks = (uniques.numel() + threads - 1) / threads;
  unique_back_map_kernel<<<blocks, threads, 0, c10::cuda::getCurrentCUDAStream()>>>(
      uniques.numel(),
      x.numel(),
      uniques.data_ptr<int64_t>(),
      cumsum_counts.data_ptr<int64_t>(),
      x.data_ptr<int64_t>(),
      back_map.data_ptr<int64_t>());

  return std::make_tuple(uniques, inverse, cumsum_counts, back_map);
}

型