pytorchである区間からある区間までの話を計算したい

Question

実現したいこと

ある配列data(mxd)とインデックス(n)が存在します.
indexはcumsum方式になっています. これはcsr_matrixのdataとindptrと同じです.

例えばdataはm=10, d=1とし[1,2,3,4,5,6,7,8,9,10]とします.
indexはi番目の要素がdataの何番目までかを示します.
例えばn=4のとき[0, 3, 5, 10]のようになっています.　これは　出力の要素が3つあり, 1番目の要素が0から2まで, 2番目の要素が3から4まで, 3番目の要素が5から9までの要素を足し合わせることを指します.
よって出力は[6(1+2+3), 9(4+5), 40(6+7+8+9+10)]となります.

これをpytorchで計算したいのですが良い計算方法はないでしょうか. よろしくお願いします.

該当のソースコード

普通に実装すると以下のようになるのですがfor文を使わず関数で綺麗に書けないでしょうか

out = torch.zeros((index.size(0)-1, data.size(1))
for i in range(index.size(0)-1):
    out[i] = torch.mean(data[index[i]:index[i+1], dim=0)

out = torch.zeros((index.size(0)-1, data.size(1))
target = torch.zeros_like(index)
for i in range(index.size(0)-1):
    target[index[i]:index[i+1]]=i
out = torch_scatter.scatter(data, target, out=out, dim=0, reduce='mean')

補足情報（FW/ツールのバージョンなど）

python 3.10
pytorch 2.0.9

Answer 1

for文を使わず関数で綺麗に書けないでしょうか

綺麗とは言えませんが、とりあえず torch.tensor_split() と torch.Tensor.expand() を使ってみました。

import torch

index = torch.tensor([0, 3, 5, 10])
data= torch.tensor([[1,2,3,4,5,6,7,8,9,10], [10,9,8,7,6,5,4,3,2,1]], dtype=torch.float)
out = torch.cat([torch.sum(i, -1).reshape(-1, i.size(0))
                 for i in torch.tensor_split(data[None, :] if data.dim() == 1 else data, index, -1)
                 if len(torch.flatten(i))]).T
print(out)

# tensor([[ 6.,  9., 40.],
#         [27., 13., 15.]])

Comments

[tanhxyx]

わかりにくくて恐縮ですが二次元の時は
data=
[[1,2,3,4,5,6,7,8,9,10]
[10,9,8,7,6,5,4,3,2,1]]
index=[0, 3, 5, 10]
の時
out=
[[6, 9, 40],
[27, 13,15]
になります

[melian]

回答を書き換えてみましたが、、よけいに複雑になってしまいました。for ループの方がよいかもしれません。

[tanhxyx]

回答ありがとうございます。いつも助かります.