Inverse of return_inverse :)

1

Hi! torch.unique has return_inverse argument. When set to True it will return an index to re-construct the original tensor from the sorted one.

sorted_and_unique, idx = torch.unique(orig_input, 0, return_inverse=True)
# recontruct the original (same as orig_input)
original = sorted_and_unique[idx]

Would there be a way to get the reverse mapping rev_idx such that orig_input[rev_idx] would result in sorted_and_unique?

The use case is that I have two tensors coords (N by 3 tensor) and features (N by 125 tensor). I would like to remove duplicates and sort coords using torch.unique as above and then use rev_idx to select the same rows in same order from feats. In other words, torch.unique offers mapping from sorted tensor to original tensor, I would like to get a mapping from original tensor to sorted tensor.

2

I think scatter_ should work:

orig_input = torch.randint(10, 15, (10,))

sorted_and_unique, idx = torch.unique(orig_input, sorted=True, return_inverse=True, dim=0)
# recontruct the original (same as orig_input)
original = sorted_and_unique[idx]

print((orig_input == original).all())
> tensor(True)

ret = torch.zeros_like(sorted_and_unique).scatter_(0, idx, orig_input)
print((ret == sorted_and_unique).all())
> tensor(True)

However, note that the result might not be deterministic as you are writing to the same location in case your other tensor has different values, which should be mapped to the same index.

1 Like