Hi guys, I am trying to implement a function that generates unique indexes for a tensor.
The rules are simple:
- assign a new unique value for each position that value equals to 2.
- keep the unique value from the previous positions if the current position is following a 2 started sequence.
- ignore value=1 positions which has no 2 a head.
Here is an example:
Input: [2, 0, 0, 2, 1, 0, 0, 2, 1, 1, 0, 0, 1, 1, 1]
Expect Output: [1, 0, 0, 2, 2, 0, 0, 3, 3, 3, 0, 0, 0, 0, 0]
I hope the whole pipeline can run within the autograd strategy.
I can successfully select and assign unique values for all values=2 right now, but it is hard for me to select values=1 and lead by 2.
>>> a = tc.Tensor([2., 0, 0, 2, 1, 0, 0, 0, 2, 1, 1, 0, 0, 1, 1, 1])
>>> B = tc.where(a == 2.0, tc.cumsum(a, -1).double(), 0.0)
>>> print(B)
tensor([2., 0., 0., 4., 0., 0., 0., 0., 7., 0., 0., 0., 0., 0., 0., 0.])