NumPyのeinsumを理解する

どのようにeinsum機能するかを正確に理解するのに苦労しています。私はドキュメントといくつかの例を見てきましたが、こだわっていないようです。

これはクラスで行った例です：

C = np.einsum("ij,jk->ki", A, B)

2つのアレイAとB

これには時間がかかると思いますA^T * Bが、よくわかりません（そのうちの1つの転置が正しく行われていますか？）。誰かが私にここで何が起こっているのか正確に説明できますeinsumか？

（注：この回答は、少し前に書いた私に関する短いブログ投稿に基づいてい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つの入力配列のラベルを付けました。の右側に->、目的の配列にラベルを付けました。

次に何が起こるかは次のとおりです。

• A1つの軸があります。ラベルを付けましたi。そして、B二つの軸を持っています。軸0をi、軸1をとラベル付けしましたj。

• 両方の入力配列でラベルを繰り返すことにより、これらの2つの軸を乗算する必要があることを示しています。言い換えれば、配列と同じように、配列の各列に配列を乗算しています。ieinsumABA[:, 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が繰り返されていることがわかります。つまり、の行Aとの列が乗算されていますB。さらに、ラベルjは出力に含まれていません。これらの製品を合計しています。ラベルiとkは出力用に保持されるため、2D配列が返されます。

この結果を、ラベルjが合計されていない配列と比較すると、さらに明確になる場合があります。下の左側には、書き込みから生じた3D配列が表示されていますnp.einsum('ij,jk->ijk', A, B)（つまり、ラベルを保持していますj）。

合計軸jは、右に示すように、予想される内積を与えます。

いくつかの演習

の感触をeinsumよくするために、添え字表記を使用しておなじみのNumPy配列演算を実装すると便利です。乗算軸と加算軸の組み合わせを含むものはすべて、を使用して記述できます einsum。

AとBを同じ長さの2つの1D配列とする。たとえば、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つの長さが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)  

numpy.einsum()直感的に理解すれば、の理念をつかむのはとても簡単です。例として、行列の乗算を含む簡単な説明から始めましょう。

を使用するにはnumpy.einsum()、いわゆる添え字文字列を引数として渡し、その後に入力配列を渡すだけです。

2つの2D配列とがAありB、行列の乗算を実行するとします。そうしたらいい：

np.einsum("ij, jk -> ik", A, B)

ここで、下付き文字列 ijは配列に対応しA、下付き文字列 jkは配列に対応していBます。また、ここで注意すべき最も重要なことは、各添え字文字列の文字数が配列の次元と一致する必要があることです。（つまり、2D配列の場合は2文字、3D配列の場合は3文字など）そして、添え字文字列の間で文字を繰り返す場合（この場合）、つまり、それらの次元に沿って合計が発生することを意味します。したがって、それらは合計削減されます。（つまり、その次元はなくなります） jein

この後の添え字文字列->は、結果の配列になります。空のままにすると、すべてが合計され、結果としてスカラー値が返されます。それ以外の場合、結果の配列は添え字文字列に応じた次元になります。この例では、になりますik。これは直感的です。なぜなら、行列の乗算では、配列の列の数が配列Aの行の数と一致する必要があることを知っているからBです（つまり、ここで起こっていることjを、添え字文字列の char を繰り返すことによってこの知識をエンコードします）。

ここでは、いくつnp.einsum()かの一般的なテンソルまたはnd-array演算を簡潔に実装する場合の使用/能力を示すいくつかの例を示します。

入力

# a vector
In [197]: vec
Out[197]: array([0, 1, 2, 3])

# an array
In [198]: A
Out[198]: 
array([[11, 12, 13, 14],
       [21, 22, 23, 24],
       [31, 32, 33, 34],
       [41, 42, 43, 44]])

# another array
In [199]: B
Out[199]: 
array([[1, 1, 1, 1],
       [2, 2, 2, 2],
       [3, 3, 3, 3],
       [4, 4, 4, 4]])

1）行列の乗算（と同様np.matmul(arr1, arr2)）

In [200]: np.einsum("ij, jk -> ik", A, B)
Out[200]: 
array([[130, 130, 130, 130],
       [230, 230, 230, 230],
       [330, 330, 330, 330],
       [430, 430, 430, 430]])

2）主対角線に沿って要素を抽出する（と同様np.diag(arr)）

In [202]: np.einsum("ii -> i", A)
Out[202]: array([11, 22, 33, 44])

3）アダマール積（2つの配列の要素ごとの積）（と同様arr1 * arr2）

In [203]: np.einsum("ij, ij -> ij", A, B)
Out[203]: 
array([[ 11,  12,  13,  14],
       [ 42,  44,  46,  48],
       [ 93,  96,  99, 102],
       [164, 168, 172, 176]])

4）要素ごとの二乗（np.square(arr)またはと同様arr ** 2）

In [210]: np.einsum("ij, ij -> ij", B, B)
Out[210]: 
array([[ 1,  1,  1,  1],
       [ 4,  4,  4,  4],
       [ 9,  9,  9,  9],
       [16, 16, 16, 16]])

5）トレース（主対角要素の合計）（と同様np.trace(arr)）

In [217]: np.einsum("ii -> ", A)
Out[217]: 110

6）行列転置（と同様np.transpose(arr)）

In [221]: np.einsum("ij -> ji", A)
Out[221]: 
array([[11, 21, 31, 41],
       [12, 22, 32, 42],
       [13, 23, 33, 43],
       [14, 24, 34, 44]])

7）（ベクトルの）外積（と同様np.outer(vec1, vec2)）

In [255]: np.einsum("i, j -> ij", vec, vec)
Out[255]: 
array([[0, 0, 0, 0],
       [0, 1, 2, 3],
       [0, 2, 4, 6],
       [0, 3, 6, 9]])

8）（ベクトルの）内積（と同様np.inner(vec1, vec2)）

In [256]: np.einsum("i, i -> ", vec, vec)
Out[256]: 14

9）軸0に沿った合計（と同様np.sum(arr, axis=0)）

In [260]: np.einsum("ij -> j", B)
Out[260]: array([10, 10, 10, 10])

10）軸1に沿った合計（と同様np.sum(arr, axis=1)）

In [261]: np.einsum("ij -> i", B)
Out[261]: array([ 4,  8, 12, 16])

11）バッチ行列乗算

In [287]: BM = np.stack((A, B), axis=0)

In [288]: BM
Out[288]: 
array([[[11, 12, 13, 14],
        [21, 22, 23, 24],
        [31, 32, 33, 34],
        [41, 42, 43, 44]],

       [[ 1,  1,  1,  1],
        [ 2,  2,  2,  2],
        [ 3,  3,  3,  3],
        [ 4,  4,  4,  4]]])

In [289]: BM.shape
Out[289]: (2, 4, 4)

# batch matrix multiply using einsum
In [292]: BMM = np.einsum("bij, bjk -> bik", BM, BM)

In [293]: BMM
Out[293]: 
array([[[1350, 1400, 1450, 1500],
        [2390, 2480, 2570, 2660],
        [3430, 3560, 3690, 3820],
        [4470, 4640, 4810, 4980]],

       [[  10,   10,   10,   10],
        [  20,   20,   20,   20],
        [  30,   30,   30,   30],
        [  40,   40,   40,   40]]])

In [294]: BMM.shape
Out[294]: (2, 4, 4)

12）軸2に沿った合計（と同様np.sum(arr, axis=2)）

In [330]: np.einsum("ijk -> ij", BM)
Out[330]: 
array([[ 50,  90, 130, 170],
       [  4,   8,  12,  16]])

13）配列のすべての要素を合計します（と同様np.sum(arr)）

In [335]: np.einsum("ijk -> ", BM)
Out[335]: 480

14）複数の軸での合計（周辺化）
（と同様np.sum(arr, axis=(axis0, axis1, axis2, axis3, axis4, axis6, axis7))）

# 8D array
In [354]: R = np.random.standard_normal((3,5,4,6,8,2,7,9))

# marginalize out axis 5 (i.e. "n" here)
In [363]: esum = np.einsum("ijklmnop -> n", R)

# marginalize out axis 5 (i.e. sum over rest of the axes)
In [364]: nsum = np.sum(R, axis=(0,1,2,3,4,6,7))

In [365]: np.allclose(esum, nsum)
Out[365]: True

15）ダブルドット積（np.sum（hadamard-product） cf. 3に類似）

In [772]: A
Out[772]: 
array([[1, 2, 3],
       [4, 2, 2],
       [2, 3, 4]])

In [773]: B
Out[773]: 
array([[1, 4, 7],
       [2, 5, 8],
       [3, 6, 9]])

In [774]: np.einsum("ij, ij -> ", A, B)
Out[774]: 124

16）2Dおよび3D配列の乗算

このような乗算は、結果を検証したい方程式の線形連立方程式（Ax = b）を解くときに非常に役立ちます。

# inputs
In [115]: A = np.random.rand(3,3)
In [116]: b = np.random.rand(3, 4, 5)

# solve for x
In [117]: x = np.linalg.solve(A, b.reshape(b.shape[0], -1)).reshape(b.shape)

# 2D and 3D array multiplication :)
In [118]: Ax = np.einsum('ij, jkl', A, x)

# indeed the same!
In [119]: np.allclose(Ax, b)
Out[119]: True

逆に、np.matmul()この検証に使用する必要がある場合、次のreshapeような同じ結果を得るには、いくつかの操作を行う必要があります。

# reshape 3D array `x` to 2D, perform matmul
# then reshape the resultant array to 3D
In [123]: Ax_matmul = np.matmul(A, x.reshape(x.shape[0], -1)).reshape(x.shape)

# indeed correct!
In [124]: np.allclose(Ax, Ax_matmul)
Out[124]: True

ボーナス：ここでより多くの数学を読んでください：Einstein-Summationそして間違いなくここに：Tensor-Notation

次元が異なるが互換性のある2つの配列を作成して、相互作用を強調します。

In [43]: A=np.arange(6).reshape(2,3)
Out[43]: 
array([[0, 1, 2],
       [3, 4, 5]])


In [44]: B=np.arange(12).reshape(3,4)
Out[44]: 
array([[ 0,  1,  2,  3],
       [ 4,  5,  6,  7],
       [ 8,  9, 10, 11]])

あなたの計算は、（2,3）と（3,4）の「ドット」（積の合計）を取り、（4,2）配列を生成します。 iの最初の点心A、最後のC; kの最後B、最初のC。 j合計によって「消費」されます。

In [45]: C=np.einsum('ij,jk->ki',A,B)
Out[45]: 
array([[20, 56],
       [23, 68],
       [26, 80],
       [29, 92]])

これは同じですnp.dot(A,B).T-転置されるのは最終出力です。

に何が起こるかを確認するにはj、C添え字をijk次のように変更します。

In [46]: np.einsum('ij,jk->ijk',A,B)
Out[46]: 
array([[[ 0,  0,  0,  0],
        [ 4,  5,  6,  7],
        [16, 18, 20, 22]],

       [[ 0,  3,  6,  9],
        [16, 20, 24, 28],
        [40, 45, 50, 55]]])

これは次の方法でも作成できます。

A[:,:,None]*B[None,:,:]

すなわち、追加kの最後までの寸法をA、およびiの正面にB（2,3,4）アレイをもたらします。

0 + 4 + 16 = 20、9 + 28 + 55 = 92など 合計しjて転置して、以前の結果を取得します。

np.sum(A[:,:,None] * B[None,:,:], axis=1).T

# C[k,i] = sum(j) A[i,j (,k) ] * B[(i,)  j,k]

私が見つけたnumpyのを：貿易のトリックを（パートII）教訓

->を使用して、出力配列の順序を示します。したがって、 'ij、i-> j'は左側（LHS）と右側（RHS）を持つと考えてください。LHSでラベルを繰り返すと、積要素が賢く計算され、合計されます。RHS（出力）側のラベルを変更することにより、入力配列に関して処理したい軸を定義できます。つまり、軸0、1などに沿って合計します。

import numpy as np

>>> a
array([[1, 1, 1],
       [2, 2, 2],
       [3, 3, 3]])
>>> b
array([[0, 1, 2],
       [3, 4, 5],
       [6, 7, 8]])
>>> d = np.einsum('ij, jk->ki', a, b)

3つの軸i、j、kがあり、jが繰り返されていることに注意してください（左側）。 i,jの行と列を表しますa。j,kのためにb。

製品を計算してj軸を揃えるには、に軸を追加する必要がありaます。（b最初の軸に沿って放送されます（？））

a[i, j, k]
   b[j, k]

>>> c = a[:,:,np.newaxis] * b
>>> c
array([[[ 0,  1,  2],
        [ 3,  4,  5],
        [ 6,  7,  8]],

       [[ 0,  2,  4],
        [ 6,  8, 10],
        [12, 14, 16]],

       [[ 0,  3,  6],
        [ 9, 12, 15],
        [18, 21, 24]]])

j右側にはないのでj、3x3x3配列の2番目の軸を合計します

>>> c = c.sum(1)
>>> c
array([[ 9, 12, 15],
       [18, 24, 30],
       [27, 36, 45]])

最後に、インデックスは右側で（アルファベット順で）逆転されるため、転置します。

>>> c.T
array([[ 9, 18, 27],
       [12, 24, 36],
       [15, 30, 45]])

>>> np.einsum('ij, jk->ki', a, b)
array([[ 9, 18, 27],
       [12, 24, 36],
       [15, 30, 45]])
>>>

アインサム方程式を読むとき、それらを命令バージョンに精神的に煮詰めることができることが最も役立つことがわかりました。

次の（面倒な）ステートメントから始めましょう。

C = np.einsum('bhwi,bhwj->bij', A, B)

最初に句読点を処理すると、2つの4文字のコンマ区切りのblob- bhwiとbhwjが矢印の前にあり、bijその後に1つの3文字のblob があることがわかります。したがって、方程式は2つのランク4テンソル入力からランク3テンソル結果を生成します。

ここで、各blobの各文字を範囲変数の名前にします。blobで文字が表示される位置は、そのテンソル内で文字が及ぶ軸のインデックスです。したがって、Cの各要素を生成する命令的合計は、Cの各インデックスに1つずつ、3つのネストされたforループで開始する必要があります。

for b in range(...):
    for i in range(...):
        for j in range(...):
            # the variables b, i and j index C in the order of their appearance in the equation
            C[b, i, j] = ...

したがって、基本的に、forCのすべての出力インデックスに対してループがあります。ここでは、範囲は未定のままにします。

次に、左側を見てみましょう- 右側に表示されない範囲変数はありますか？私たちの場合-はい、そして。このような変数ごとに内部のネストされたループを追加します。hwfor

for b in range(...):
    for i in range(...):
        for j in range(...):
            C[b, i, j] = 0
            for h in range(...):
                for w in range(...):
                    ...

最も内側のループ内で、すべてのインデックスが定義されたので、実際の合計を書き込むことができ、変換が完了します。

# three nested for-loops that index the elements of C
for b in range(...):
    for i in range(...):
        for j in range(...):

            # prepare to sum
            C[b, i, j] = 0

            # two nested for-loops for the two indexes that don't appear on the right-hand side
            for h in range(...):
                for w in range(...):
                    # Sum! Compare the statement below with the original einsum formula
                    # 'bhwi,bhwj->bij'

                    C[b, i, j] += A[b, h, w, i] * B[b, h, w, j]

これまでにコードをたどることができたなら、おめでとうございます！これで、アインサム方程式を読み取ることができます。特に、元のアインサム式が上記のスニペットの最終的な合計ステートメントにどのようにマップされるかに注意してください。forループと範囲の境界は綿毛であり、その最後のステートメントは何が起こっているのかを理解するために本当に必要なすべてです。

完全を期すために、各範囲変数の範囲を決定する方法を見てみましょう。さて、各変数の範囲は、単にそれがインデックスを付ける次元の長さです。明らかに、変数が1つ以上のテンソルで複数の次元にインデックスを付ける場合、それらの各次元の長さは等しくなければなりません。上記のコードと完全な範囲は次のとおりです。

# C's shape is determined by the shapes of the inputs
# b indexes both A and B, so its range can come from either A.shape or B.shape
# i indexes only A, so its range can only come from A.shape, the same is true for j and B
assert A.shape[0] == B.shape[0]
assert A.shape[1] == B.shape[1]
assert A.shape[2] == B.shape[2]
C = np.zeros((A.shape[0], A.shape[3], B.shape[3]))
for b in range(A.shape[0]): # b indexes both A and B, or B.shape[0], which must be the same
    for i in range(A.shape[3]):
        for j in range(B.shape[3]):
            # h and w can come from either A or B
            for h in range(A.shape[1]):
                for w in range(A.shape[2]):
                    C[b, i, j] += A[b, h, w, i] * B[b, h, w, j]