どうやってnp.einsum
仕事?
配列A
とが与えられた場合B
、それらの行列乗算とそれに続く転置は を使用して計算されます(A @ B).T
。または、同等に、以下を使用して計算されます。
np.einsum("ij, jk -> ki", A, B)
ベストアンサー1
(注:この回答は短いブログ投稿einsum
少し前に書いたものです。
何einsum
をするのですか?
2 つの多次元配列があり、 があるとしますA
。B
ここで、次の操作を実行したいとします...
-
A
特定の方法で掛け合わせてB
新しい製品群を作り、そしておそらく - この新しい配列を特定の軸に沿って合計し、その後おそらく
- 新しい配列の軸を特定の順序で転置します。
を使用すると、、、などeinsum
の NumPy 関数の組み合わせよりも高速かつメモリ効率よくこれを実行できる可能性が高くなります。multiply
sum
transpose
どのようにeinsum
機能しますか?
ここに簡単な(しかし完全に簡単ではない)例を示します。次の 2 つの配列を考えてみましょう。
A = np.array([0, 1, 2])
B = np.array([[ 0, 1, 2, 3],
[ 4, 5, 6, 7],
[ 8, 9, 10, 11]])
A
要素ごとに乗算しB
、新しい配列の行に沿って合計します。「通常の」NumPy では次のように記述します。
>>> (A[:, np.newaxis] * B).sum(axis=1)
array([ 0, 22, 76])
したがって、ここでは、インデックス操作によってA
2 つの配列の最初の軸が整列し、乗算がブロードキャストできるようになります。次に、積の配列の行が合計され、答えが返されます。
代わりにを使用したい場合はeinsum
、次のように記述できます。
>>> np.einsum('i,ij->i', A, B)
array([ 0, 22, 76])
ここで鍵となるのは署名文字列'i,ij->i'
であり、少し説明が必要です。これは 2 つの部分に分けて考えることができます。左側 ( の左側->
) には、2 つの入力配列にラベルを付けています。 の右側には->
、最終的に取得する配列にラベルを付けています。
次に何が起こるかを説明します。
A
には 1 つの軸があり、それを とラベル付けしましたi
。 にB
は 2 つの軸があり、軸 0 をi
、軸 1 を とラベル付けしましたj
。両方の入力配列でラベルを繰り返すことで、これら 2 つの軸を掛け合わせる
i
ように指示しています。つまり、と同様に、配列 を配列の各列と掛け合わせています。einsum
A
B
A[:, np.newaxis] * B
希望する出力に がラベルとして表示されないことに注意してください
j
。 を使用しましたi
(1D 配列にしたい)。ラベルを省略することで、この軸に沿って合計するeinsum
ように指示しています。 つまり、 と同様に、積の行を合計しています。.sum(axis=1)
基本的に、 を使用するのに必要なのはこれだけですeinsum
。 について少し試してみると役に立ちます。出力に両方のラベルを残した場合、'i,ij->ij'
は積の 2D 配列を返します ( と同じA[:, np.newaxis] * B
)。出力ラベルを指定しない場合、'i,ij->
は単一の数値を返します ( を実行した場合と同じ(A[:, np.newaxis] * B).sum()
)。
ただし、この方法の優れた点はeinsum
、最初に一時的な積の配列を作成せず、積を合計していくだけであることです。これにより、メモリ使用量を大幅に節約できます。
少し大きめの例
ドット積を説明するために、次の 2 つの新しい配列を示します。
A = array([[1, 1, 1],
[2, 2, 2],
[5, 5, 5]])
B = array([[0, 1, 0],
[1, 1, 0],
[1, 1, 1]])
を使用してドット積を計算しますnp.einsum('ij,jk->ik', A, B)
。次の図は、 と のラベルと、関数から取得した出力配列を示していA
ますB
。
ラベルが繰り返されていることがわかります。これは、 の行と のj
列を掛け合わせていることを意味します。さらに、ラベルは出力に含まれません。これらの積を合計しています。ラベルと は出力用に保持されるため、2D 配列が返されます。A
B
j
i
k
j
この結果を、ラベルを合計しない配列と比較すると、さらに明確になるかもしれません。以下、左側に、書き込みの結果である 3D 配列が表示されますnp.einsum('ij,jk->ijk', A, B)
(つまり、ラベルは保持されていますj
)。
軸を合計するとj
、右側に示すように、予想されるドット積が得られます。
いくつかの練習
の感覚をつかむにはeinsum
、添え字表記を使用して使い慣れた NumPy 配列操作を実装すると便利です。軸の乗算と合計の組み合わせを含むものはすべて、 を使用して記述できますeinsum
。
A と B を同じ長さの 2 つの 1 次元配列とします。たとえば、A = np.arange(10)
およびですB = np.arange(5, 15)
。
の合計は
A
次のように書けます。np.einsum('i->', A)
要素ごとの乗算は
A * B
次のように表すことができます。np.einsum('i,i->i', A, B)
内積またはドット積、
np.inner(A, B)
またははnp.dot(A, B)
、次のように表すことができます。np.einsum('i,i->', A, B) # or just use 'i,i'
外積は次の
np.outer(A, B)
ように表すことができます。np.einsum('i,j->ij', A, B)
2D 配列C
およびの場合D
、軸の長さに互換性がある (両方の軸の長さが同じか、いずれかの軸の長さが 1) と仮定すると、次の例が考えられます。
C
(主対角線の和)の軌跡は次np.trace(C)
のように表すことができます。np.einsum('ii', C)
の要素ごとの乗算
C
と の転置、 はD
次C * D.T
のように表すことができます。np.einsum('ij,ji->ij', C, D)
の各要素を
C
配列で乗算してD
(4D 配列を作成する)、C[:, :, None, None] * D
は次のように記述できます。np.einsum('ij,kl->ijkl', C, D)