らんだむな記憶

blogというものを体験してみようか!的なー

einsum (1)

ここいらで {numpy,tf,torch}.einsum の動きを見ておきたい。手軽さの観点で numpy で見る。

>>> a = np.arange(1,3)
>>> a
array([1, 2])

の時、

>>> np.einsum('i,i', a,a)
5

である。これは、-> での行き先がないので、結果がスカラになるということで、

\begin{align*}
c = a_i a_i
\end{align*}

つまり

\begin{align*}
c = \sum_i a_i a_i = 1 \times 1 + 2 \times 2 = 5
\end{align*}

でドット積となる。次に、

>>> a = np.arange(1,5).reshape(2,2)
>>> a
array([[1, 2],
       [3, 4]])
>>> b = np.arange(-5,-1).reshape(2,2)
>>> b
array([[-5, -4],
       [-3, -2]])

を用意する。

>>> np.einsum('ij,ij', a, b)
-30

である。これは、


\begin{align*}
c = a_{ij} b_{ij} = \sum_{i,j} a_{ij} b_{ij} = 1 \times (-5) + 2 \times (-4) + 3 \times (-3) + 4 \times (-2) = -30
\end{align*}

ということである。

>>> np.einsum('ij,ij->i', a, b)
array([-13, -17])
>>> np.einsum('ij,ij->j', a, b)
array([-14, -16])

はそれぞれ・・・ちょっとオリジナルのアインシュタインの規約通りではないと思うのだが、

\begin{align*}
(c_i) = (\sum_j a_{ij} b_{ij}) = \begin{pmatrix} 1 \times (-5) + 2 \times (-4) \\ 3 \times (-3) + 4 \times (-2) \end{pmatrix} = \begin{pmatrix} -13 \\ -17 \end{pmatrix}
\end{align*}

\begin{align*}
(c_j) = (\sum_j a_{ij} b_{ij}) = (1 \times (-5) + 3 \times (-3)\quad 2 \times (-4) + 4 \times (-2)) = (-14\ \; -16)
\end{align*}

となる。1 階のテンソル (ベクトル) の場合、すべて “横ベクトル” 表記になるので、分かりにくいが丁寧に数式として考えると、なるほどという感じではある。本当は・・・

>>> np.einsum('ij,ij->i', a, b).reshape(2, 1)
array([[-13],
       [-17]])

が得られると数学的な感覚と符合して分かりやすいのだが・・・*1

オリジナルのアインシュタインの規約の考え方では、“右辺” だけで閉じていて、右辺だけ見ればどの添字について和をとるか分かるが、numpy.einsum などでは、“左辺” の添字を見て、どの添字が生き残るか、つまりどの添字については和をとらないかを見なければいけない状態にある。

ちょっとまとめると、einsum の分かりにくさは

  • オリジナルのアインシュタインの規約から拡張されている*2
  • 具体的な例としては低階のテンソル、つまりスカラとベクトルが扱いやすいが、ベクトルの縦横の概念がなくなってすべて “横ベクトル” 表記で出力されるので、数学的な慣習で考えすぎると和をとる方向が頭の中で混乱する

部分にあるような気がする。

*1:とは言え、任意階のテンソルを考えるにあたって、ベクトルや行列の表記上の慣習に合わせる必要はないかもしれない。そもそも shape による次元が (2,) ではなく (2, 1) になるため、2 階のテンソルになってしまうので都合が悪い・・・。

*2:同じ項で添字が重なる場合はその添字について和を取る・・・はずなのに、重なる場合でも和をとらない指定が平気でできる。