行列をより効率的に索引付けする（tf.gatherを避ける）？

Question

Bの（次元N×Kの）2つの行列AおよびインデックスiA =（a1、a2 ... aM）のリストが与えられ、BのiB =（b1、b2 ... bM）以下を実行する：

lp = 0 
for a in iA: 
    for b in iB: 
     lp += np.sum(A[a,] * B[b,]) 

in Tensorflow。インデックスのリストには繰り返しがあるので、同じ行を複数回描画します。

私の現在の実装では、次のようになります（私はtf.gather使用していたためと思われる）

lp = tf.reduce_sum(tf.multiply(tf.gather(A, iA), tf.gather(B, iB)), 1) 

はしかし、勾配は計算するのは非常に遅いです。 iAは昇順でソートされていると仮定できます（したがって、特定の行A [a、]は 'a'が変更されるまで再利用できます）。 これを行うより良い方法はありますか？どんな助けもありがとうございます。

編集：AとBはここではtfです。変数と繰り返しごとに変わります。事前計算はオプションではありません。しかしながら、iA及びiBは定数であり、変化しない。

編集：M >> N.Kのサイズは実際問題ではありません。

Answer 1

インデックスの密度によっては、単純にカウントの配列を事前計算することをお勧めします。あなたは、単にこれに乗じ、合計することにより、インデックス付きの合計を計算することができ

counts = np.zeros_like(A) 
np.add.at(counts, iA, 1) 
np.add.at(counts, iB, 1) 

だろうnumpyのでは：

lp = np.sum(counts * A * B) 

をしかし、iAとiBは両方とも非常に疎である場合、これは非常にいくつかを与えますパフォーマンスオーバーヘッドiAまたはiBが固定されていない場合でも同じことが適用されますが、反復ごとに変更されます。