Substitute values in a vector (PyTorch)

I have a tensor of N unique target labels, randomly selected from [0,R], where N<R (i.e., my target vector can have any length, but only contains N unique labels.). I would like to transform the labels to [0,N]. Is there a function available for this target transform? e.g. input vector: [12, 6, 4, 5, 3, 12, 4] → transformed vector : [4, 3, 1, 2, 0, 4, 1]

My attempt:

I have implemented the following snippet, which works as expected, but might not be the most glorious implementation:

import torch


def my_transform(vec):
    t_ = torch.unique(vec)
    return torch.cat(list(map(lambda x: (t_ == x).nonzero(as_tuple=True)[0], vec)))


t = torch.Tensor([12, 6, 4, 5, 3, 12, 4])
print(my_transform(t))

Solution 1:

You're looking for searchsorted

import torch

t = torch.Tensor([12, 6, 4, 5, 3, 12, 4])
transformed = torch.searchsorted(t.unique(),t)
# tensor([4, 3, 1, 2, 0, 4, 1])