Pytorch中计算张量不等大块总和的有效方法

Question

我想知道是否有办法计算pytorch中张量不等大小的以下总和

import torch
import numpy as np

x = torch.rand(1000,100)
y = np.unique(np.random.choice(1000,10)

这里我有一个尺寸为张量x（1000,10），我想计算沿第一轴的夹头总和。这些块沿第一轴分开，y表示每个块的结束线。它们通常具有不相等的大小。例如，我可以使用以下for循环

执行此操作

cum_pos_lik = torch.FloatTensor(y.size, 100)
y = np.append(0, y)
for i in range(1, y.size):
    cum_pos_lik[i-1, :] = x[y[i-1]:y[i], :].sum(0)

但我需要这样才能更快地完成我的应用程序。显然，每个块的总和可以并行化。我想知道在pytorch中是否有一个简单的方法。如果提供一种在numpy中矢量化它的方法也会很好。

谢谢！