How to shift values in tensor
I have tensor T of shape [batch_size, A] with values and tensor S of shape [batch_size] with shift parameters.
I would like to shift values in T[b] by S[b] positions to the right, the last S[b] elements of T[b] should be dropped and new elements should be set to 0.
So basically want to do something like:
for i in range(batch_size): T[i] = zeros[:S[i]] + T[i, :A-S[i]]
For: T = [[1, 2, 3], [4, 5, 6]] S = [1, 2] Return: T' = [[0, 1, 2], [0, 0, 4]]
Is there some easy way to do it?