Generating indices for repeating values in a tensor

1

Hi,
I have a very simple question.
If i have an tensor array which is like [1,1,1,1,2,2,3,3,3,2,2]
is there any function or method to count the repetition and generate another tensor which is: [0,1,2,3,0,1,0,1,2,0,1]

2

The simple way would be to use a window of two values in a loop:

x = [1,1,1,1,2,2,3,3,3,2,2]
res = [0]
for idx, (a, b) in enumerate(zip(x[:-1], x[1:])):
    if a == b:
        tmp = res[idx] + 1
        res.append(tmp)
    else:
        res.append(0)

I’m not sure, if any PyTorch/numpy method might speed it up, as e.g. unique seems to be fruitless here, since you are repeating values.