如何向量化总和聚合？

Question

我试图将这段 Python（或经过一些小改动后的 Matlab）代码矢量化，这是有向图的总和聚合

for j in range(batchSize):
    for i in range(2*nEdges[j]):
        localSum[j,receivers[j,i],0:2]+=(localFeature[j,i,0:2])

我知道如何在 = 时删除其中一个循环，但我没能找到在 += 情况下执行此操作的方法。做同样的技巧的困难在于 receivers[j,i] 对于不同的 j 和 i 可能是相同的值。

有人知道在这种情况下如何进行矢量化吗？

非常感谢:)

Answer 1

终于找到了。 torch_scatter 包提供此功能（用于 pytorch 张量）。嵌套循环可以重写为

localSum = torch_scatter.scatter(localFeature,receivers,dim=1,reduce='sum')

How to vectorize the sum aggregation?