如何对总和聚合进行矢量化

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])

我知道当它是=时如何丢弃其中一个循环，但我没能找到一种方法来处理这个+=的情况。实现相同技巧的困难在于，对于不同的j和is，receivers[j,i]可能是相同的值。

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

非常感谢：(

Answer 1

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

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