pytorchの行列同士の掛け算について

やりたいこと

3x3行列であるxと、1次元行列の配列であるyが与えられたとき、x*yを綺麗に計算したいです。（説明がわかりにくくてすみません。）
例えば

# x.shape = ([3, 3])
x = torch.tensor([
    [  1,   2, 1], 
    [1.5, 2.5, 1],
    [  0,   0, 1]],
dtype=torch.float)

# y.shape = torch.Size([2, 2, 3])、yは3次元である想定。
y = torch.tensor([
[[1, 1, 1], 
 [1, 2, 1]], 
[[2, 1, 1], 
 [2, 2, 1]]
],
dtype=torch.float)

が与えられたとき、

[  1,   2, 1]    [1]    [4]
[1.5, 2.5, 1] ×  [1]  = [5]
[  0,   0, 1]    [1]    [1]

[  1,   2, 1]    [1]    [6]
[1.5, 2.5, 1] ×  [2]  = [7.5]
[  0,   0, 1]    [1]    [1]

...

のように、xとyの各要素の積を計算して、
after_y = torch.tensor([[[4,5,1], [6, 7.5, 1]], [[5, 6.5, 1],[7, 9, 1]]])
のように、元のyと同じ次元で答えを返すようにしたいです。

現在のコード

書き方が分からず、とりあえずやりたいことをfor文で書きました。。

after_y = torch.zeros(y.shape)

for i in range(after_y.shape[0]):
    for j in range(after_y.shape[1]):
        after_y[i,j] = torch.matmul(x, y[i,j])

何の関数を使ってどう直したらいいのかが分からず、
助けていただけると嬉しいです。

行動規範の内容に同意します

回答1件

ベストアンサー

python
1>>> torch.mm(x, y.reshape(-1, 3).T).T.reshape(y.shape)
2tensor([[[4.0000, 5.0000, 1.0000],
3         [6.0000, 7.5000, 1.0000]],
4
5        [[5.0000, 6.5000, 1.0000],
6         [7.0000, 9.0000, 1.0000]]])