举个例子，我有两个张量：A = [1;2;3;4;5;6;7] 和 B = [2;3;2]。我的想法是基于 B 来降维 A，使得 B 的值表示如何对 A 的值进行求和——具体来说，B = [2;3;2] 意味着降维后的 A 将是前两个值的和，接下来的三个值的和，以及最后两个值的和：A' = [(1+2);(3+4+5);(6+7)]。显然，B 的总和应始终等于 A 的长度。我希望尽可能高效地实现这一点——最好使用 pytorch/python 中的特定函数或矩阵操作。谢谢！

回答：

这是解决方案。

• 首先，我们创建一个与 A 相同大小的索引数组 B_idx。
• 然后，根据索引 B_idx 使用 index_add_ 来累积（加总）A 中的所有元素。

A = torch.arange(1, 8) B = torch.tensor([2, 3, 2])B_idx = [idx.repeat(times) for idx, times in zip(torch.arange(len(B)), B)]B_idx = torch.cat(B_idx) # tensor([0, 0, 1, 1, 1, 2, 2])A_sum = torch.zeros_like(B)A_sum.index_add_(dim=0, index=B_idx, source=A)print(A_sum) # tensor([ 3, 12, 13])