Pythonの効率的なテンソル収縮

Question

私はテンソル（ndarrayオブジェクト）のリストLをいくつか持っています。私は接続のグラフに従ってこれらのインデックスを契約する必要があります。

接続は「テンソルL[n]のJ番目のインデックスを持つテンソルL[m]のI番目のインデックスを委託意味フォーム((m,i),(n,j))でタプルのリストでエンコードされている。

どのようにすることができます私は非自明な接続グラフを扱いますか？最初の問題は、インデックスのペアを契約すると、リストに属していない新しいテンソルが結果になるということです。すべてのテンソルのすべてのインデックスに識別子）、収縮を実行するための命令を選択できるという問題があります。いくつかの選択肢計算の途中で不必要に巨大な獣を生み出します（最終結果が小さい場合でも）。提案？

Answer 1

メモリの問題とは別に、einsumへの1回の呼び出しで収縮を行うことができると思いますが、いくつかの前処理が必要です。私は完全にあなたが "のインデックスのペアを契約して、結果はリストに属していない新しいテンソルですL"が意味するものではないが、私は1つのステップで収縮を行うことが正確に解決すると思いますこの問題。私は、einsumの数値添字構文の代替を使用することをお勧め

：

einsum(op0, sublist0, op1, sublist1, ..., [sublistout]) 

だから、あなたが何をする必要があるか、整数配列中に収縮するインデックスを符号化しています。最初に、一意のインデックスの範囲を最初に設定し、別のコピーをsublistoutとして保持する必要があります。次に、接続グラフを繰り返し、必要な箇所で同じインデックスに契約インデックスを設定すると同時に、契約インデックスをsublistoutから削除する必要があります。

import numpy as np 

def contract_all(tensors,conns): 
    ''' 
    Contract the tensors inside the list tensors 
    according to the connectivities in conns 

    Example input: 
    tensors = [np.random.rand(2,3),np.random.rand(3,4,5),np.random.rand(3,4)] 
    conns = [((0,1),(2,0)), ((1,1),(2,1))] 
    returned shape in this case is (2,3,5) 
    ''' 

    ndims = [t.ndim for t in tensors] 
    totdims = sum(ndims) 
    dims0 = np.arange(totdims) 
    # keep track of sublistout throughout 
    sublistout = set(dims0.tolist()) 
    # cut up the index array according to tensors 
    # (throw away empty list at the end) 
    inds = np.split(dims0,np.cumsum(ndims))[:-1] 
    # we also need to convert to a list, otherwise einsum chokes 
    inds = [ind.tolist() for ind in inds] 

    # if there were no contractions, we'd call 
    # np.einsum(*zip(tensors,inds),sublistout) 

    # instead we need to loop over the connectivity graph 
    # and manipulate the indices 
    for (m,i),(n,j) in conns: 
     # tensors[m][i] contracted with tensors[n][j] 

     # remove the old indices from sublistout which is a set 
     sublistout -= {inds[m][i],inds[n][j]} 

     # contract the indices 
     inds[n][j] = inds[m][i] 

    # zip and flatten the tensors and indices 
    args = [subarg for arg in zip(tensors,inds) for subarg in arg] 

    # assuming there are no multiple contractions, we're done here 
    return np.einsum(*args,sublistout) 

簡単な例：私たちはすべての重複を処理する必要があるため、複数の収縮がある場合に

>>> tensors = [np.random.rand(2,3), np.random.rand(4,3)] 
>>> conns = [((0,1),(1,1))] 
>>> contract_all(tensors,conns) 
array([[ 1.51970003, 1.06482209, 1.61478989, 1.86329518], 
     [ 1.16334367, 0.60125945, 1.00275992, 1.43578448]]) 
>>> np.einsum('ij,kj',tensors[0],tensors[1]) 
array([[ 1.51970003, 1.06482209, 1.61478989, 1.86329518], 
     [ 1.16334367, 0.60125945, 1.00275992, 1.43578448]]) 

は、ループ内の物流は、もう少し複雑になります。しかし、論理は同じです。さらに、上記は対応する指標がに収縮することを確実にするためのチェックが明らかに欠落している。

私は、デフォルトのsublistoutを指定する必要がないことを認識しましたが、einsumはその順序を使用します。その変数をコードに残すことにしました。なぜなら、わかりやすい出力インデックスの順序が必要な場合は、その変数を適切に処理する必要があり、便利な場合があるからです。 （今は削除されたコメントで@hpauljで述べたように）

---

は収縮順序の最適化のためとして、あなたは、バージョン1.12のようnp.einsumに内部の最適化を行うことができます。このバージョンでは、というキーワード引数をnp.einsumに導入しました。これにより、計算時間をメモリコストで削減する縮小順序を選択できます。 'greedy'または'optimal'をoptimizeキーワードとして渡すと、numpyは次元のサイズの大まかな減少順に縮小順を選択します。（オンラインドキュメントが行く限り、help()幸い作品）明らかに文書化されていないから来optimizeキーワードに利用できる

オプション機能np.einsum_pathは：

einsum_path(subscripts, *operands, optimize='greedy') 

Evaluates the lowest cost contraction order for an einsum expression by 
considering the creation of intermediate arrays. 

np.einsum_pathから出力収縮パスとしても使用することができますoptimize引数の入力はnp.einsumです。あなたの質問では、あまりにも多くのメモリが使用されていることを心配していました。したがって、デフォルトでは、最適化が行われていないことが考えられます（ランタイムが長くなり、メモリフットプリントが小さくなる可能性があります）。