根据张量元素在pytorch中的位置向其添加值

Question

在pytorch中，我想根据元素的位置向张量中的元素添加值。例如，考虑，

Input = torch.tensor([1,2,3,4,5,6,7,8,9,0,1,2,3,4,5,6,7,8,9,0])

在输入数组的几个偏移量之间Offsets = [0,5,10,15,20]，我想添加不同的值ValuesToAdd = [10,100,1000,10000]

我希望输出是

Output = torch.tensor([11,12,13,14,15,106,107,108,109,100,1001,1002,1003,1004,1005,10006,10007,10008,10009,10000])

这里，在输入数组中的索引Offsets[i]和Offsets[i+1]之间添加ValuesToAdd[i]。例如，对于输入数组中的索引10,11,12,13和14（Offsets[2] = 10到Offsets[3]=15），添加了1000（ValuesToAdd[2]）

我怎样才能做到这一点？我不是在偏移量数组上循环，而是在寻找一种更有效的方法

Answer 1

您可以使用torch.repeat_interleave

Offsets = torch.tensor(Offsets)
shifts = Offsets[1:] - Offsets[:-1]
output = Input.clone()
output[Offsets[0]:Offsets[-1]] += torch.tensor(ValuesToAdd).repeat_interleave(shifts)
print(torch.all(output == Output))
# True