PyTorch：作为张量运算的动态规划

Question

是否可以通过Tensor运算得到如下循环？

a = torch.Tensor([1, 0, 0, 0])
b = torch.Tensor([1, 2, 3, 4])

for i in range(1, a.shape[0]):
    a[i] = b[i] + a[i-1]

print(a) # [1, 3, 6, 10]

该操作取决于 a 中的先前值和在途中计算的值（以动态编程方式）。

是否可以通过张量运算来实现这种顺序计算？

Answer 1

您可以通过 cumulative sum:

b.cumsum(0)

PyTorch: Dynamic Programming as Tensor Operation