numpyの比較的よく使う関数に、dot積があります。
スカラーとベクトルを渡せばベクトルをスカラー倍してくれて、ベクトル同士なら内積を取ってくれ、
行列を二つ渡せばそれらの行列積を戻してくれるとても便利な関数です。
(本当は行列積は、 np.matmul を使ったほうがいいらしい。)
さて、そのnp.dot ですが、行列よりもより高次元の配列についても定義されていることを最近知りました。
ドキュメント : numpy.dot
二つの多次元配列$a$と$b$に対して、$a$の一番最後の次元の長さと、$b$の最後から２番目の次元の長さが等しい時に、
np.dot(a, b)を計算することができます。
i*j行列と、j*k行列に積が定義されて、その積がi*k行列になるのと似ています。
具体的な挙動について、コード動かしてみていきましょう。
まず、サンプルとなるデータを作ります。
aの最後の次元の要素数と、bの最後から2番目の次元の要素数は 5で揃えましたが、
それ以外の次元の要素数はバラバラにして、結果と比較しやすいようにしました。
import numpy as np
a = np.random.randn(2, 3, 4, 5).round(2)
b = np.random.randn(6, 7, 5, 8).round(2)
print(a.shape, b.shape)
# (2, 3, 4, 5) (6, 7, 5, 8)
このdot積を取って、shapeをみてみます。
c = np.dot(a, b)
print(c.shape)
# (2, 3, 4, 6, 7, 8)
5は消えましたが残りの数はそのまま残りましたね。
さて、結果の$c$の各要素の値ですが、次のように計算されたものが入っています。
$$
c[i, j, k, m, n, o] = \sum_{l} a[i, j, k, l] \cdot b[m, n, l, o].
$$
別の書き方をすればこうです。
$$
c[i, j, k, m, n, o] = np.dot(a[i, j, k, :], b[m, n, :, o]).
$$
一応確認しておきましょう。
c[1, 2, 3, 4, 5, 6] == np.dot(a[1, 2, 3, :], b[4, 5, :, 6])
# True