I’m trying to implement the following function as efficiently as possible:
Given I have a sorted 1D tensor of positive integers, I want to retrieve the index of the first occurrence of each value. In case any integer between 0 and N-1 is missing, I want to return a -1. Thus, the function would return a 1D tensor of N elements.
Is there an efficient way of doing this better than just a for loop over all values?
x = torch.tensor([2, 3, 3, 4, 6, 6, 6, 8, 8], dtype=torch.long)
y = arg_first_ocurrence(x, N=10)
tensor([-1, -1, 0, 1, 3, -1, 4, -1, 7, -1])
I’m open to solutions that would involve coding a C++/CUDA extension.