Pythonのeinsum_path()でeinsum式の最低コストの縮約順序を評価する方法

はじめに

einsum式における最低コスト（最適）な縮約順序を求めるには、Pythonのnumpy.einsum_path()メソッドを使用します。第1引数のsubscripts（添字）は総和を行うための添字を指定し、第2引数のoperands（オペランド）には演算対象となる配列を渡します。

アインシュタインの縮約記法（Einstein summation convention）を用いることで、多次元配列に対する多くの一般的な線形代数操作を簡潔に表現できます。暗黙モードでは、einsumはこれらの値をそのまま計算します。

一方、明示モードでは、特定の添字ラベルへの総和を無効化または強制できるため、古典的なアインシュタイン縮約とはみなされないような配列操作も柔軟に計算できます。

得られるパスは、入力の縮約式のどの項を先に縮約すべきかを示しています。この縮約結果は縮約リストの末尾に追加され、すべての中間縮約が完了するまでリストを繰り返し処理していきます。

手順

まず、必要なライブラリをインポートします。

import numpy as np

次に、テンソルを作成します。

p = np.random.rand(2, 2)
q = np.random.rand(2, 5)
r = np.random.rand(5, 2)

einsum式の最低コストの縮約順序を取得するには、numpy.einsum_path()メソッドを使用します。

path_info = np.einsum_path('ij,jk,kl->il', p, q, r, optimize='greedy')

パス情報を表示します。

print(path_info[0])
print(path_info[1])

実行例

import numpy as np
np.random.seed(123)

# テンソル
p = np.random.rand(2, 2)
q = np.random.rand(2, 5)
r = np.random.rand(5, 2)

# einsum_path()メソッドで最低コストの縮約順序を取得
path_info = np.einsum_path('ij,jk,kl->il', p, q, r, optimize='greedy')

# パス情報を表示
print(path_info[0])
print(path_info[1])

出力

['einsum_path', (1, 2), (0, 1)]
  Complete contraction:  ij,jk,kl->il
         Naive scaling:  4
     Optimized scaling:  3
      Naive FLOP count:  1.200e+02
  Optimized FLOP count:  5.700e+01
   Theoretical speedup:  2.105
  Largest intermediate:  4.000e+00 elements
--------------------------------------------------------------------------
scaling                 current                   remaining
--------------------------------------------------------------------------
   3                 kl,jk->jl                  ij,jl->il
   3               jl,ij->il                     il->il

出力結果の読み方

パス情報の最初の要素['einsum_path', (1, 2), (0, 1)]は、まず入力テンソルのうち2番目と3番目（インデックス1と2）にあたるqとrを縮約し、続いてその結果とpを縮約することを意味しています。

また、ナイーブな計算量（Naive scaling: 4）に対して、最適化後の計算量（Optimized scaling: 3）が抑えられており、理論上は約2.1倍の高速化が期待できることがわかります。optimize='greedy'を指定することで、貪欲法によるコスト最小化の探索が行われます。