More than 5 years have passed since last update.

numpyにおける tensordotとeinsumの関係

Posted at 2019-02-25

はじめに

個人的な備忘録を兼ねて簡潔に書きます。
誤った理解をしている可能性もあります。
その際は、コメントで指摘して下さい。

einsumの方がtensordotより一般的にかけるものと理解しています。

詳しい内容は下記の公式リファレンスやブログを参考にして下さい。

numpy.tensordot(A, B, axes)

A、Bは多次元配列 axesは掛け合わす軸（整数型）を指定する。

個人的には、tensordotは何をやっているかがわかりづらかったです。
einsumとの関係を調べ、tensordotの処理を理解できたので、書き留めます。

多くの tensordotは以下のeinsumで同じような処理を実現できます

numpy.einsum(subscripts, A, B)

subscripts は文字列型で変換前の添字と変換後の添字を指定します。
'{Aの添字}, {Bの添字}->{最終的な添字}'といった具合に指定します。

例えば、

C = numpy.einsum('ijk, ijc->kc', A, B)

とすれば
$C_{kc} = \sum_{i, j} A_{ijk} B_{ijc}$
の計算を行います。
簡単にいうと iとjについて総和をとり、kとbが残るという感じです。

例に挙げた einsumと同じ処理をするtensordotは

C_hat = numpy.tensordot(A, B, ([0, 1], [0, 1]))

となります。

tensordotでは総和をとる軸を指定するので、以上のような形になります。

最終的な次元（Cの次元）はAとBで残った次元を合わせたものになります。

A.shape = (2, 3, 3), B.shape = (2, 3, 4)
ならば
C.shape = (3, 4)
です。

一言でまとめると
tensordot では消す（足し合わす）軸を axesで指定しているだけです。