Numpyのeinsumを使って３次元配列の掛け算

知り合いに簡単に3次元配列の掛け算をする方法を聞かれ、調べてみたらNumpyのライブラリにアインシュタインの縮約記号を用いたeinsumというのがありこれを使えば楽にできそうなのでやってみました。

１ｘ２ｘ５の3次元行列と２ｘ３ｘ５の三次元行列をまず適当に2つ作ります。

A = 
\left(
\begin{matrix}
a_0 & a_1
\end{matrix}
\right)
, a_0 = 
\left[
\begin{matrix}
0 & 1 & 1 & 0 & 0
\end{matrix}
\right]
, a_1 = 
\left[
\begin{matrix}
0 & 1 & 0 & 0 & 1
\end{matrix}
\right]

W = 
\left(
\begin{matrix}
w_0 & w_1 & w_2 \\
w_3 & w_4 & w_5
\end{matrix}
\right)
, w_0 = 
\left[
\begin{matrix}
0 & 1 & 2 & 3 & 4
\end{matrix}
\right]
, w_1 = 
\left[
\begin{matrix}
5 & 6 & 7 & 8 & 9
\end{matrix}
\right]
...

これをnumpyを使って作成すると、

A = np.array([[[0,1,1,0,0],
               [0,1,0,0,1]]])

W = np.arange(30).reshape(2,3,5)

で出来ます。

このAの行列の形状（１，２，５）をアインシュタインの縮約記号で（i, j, k)で表します。
同じようにWの行列の形状（２，３，５）を（j, l, k)で表します。

この２つの行列を掛け算すると、（１，３，５）の行列になるはずなのでeinsumを使った式はこうなります。

R = np.einsum("ijk,jlk -> ilk",A,W)

print(R)
print(R.shape)

実行するとRはこうなります。

[[[ 0 17  2  0 19]
  [ 0 27  7  0 24]
  [ 0 37 12  0 29]]]

(1,3,5)

無事、１ｘ３ｘ５の３次元行列の結果になりました。

参考（英語）：

• https://stackoverflow.com/questions/26089893/understanding-numpys-einsum