Indizieren von Matrizen effizienter (Vermeidung von Gattern)?

Question

Gegeben zwei Matrizen A und B (der Dimensionen N x K) und eine Liste der Indizes iA = (a1, a2 ... aM) für A und iB = (b1, b2 ... bM) für B, brauche ich Folgendes zu tun:

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

in Tensorflow. Die Liste der Indizes hat Wiederholungen, daher zeichnen wir dieselbe Zeile mehrmals.

Meine aktuelle Implementierung sieht wie folgt aus:

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

Allerdings sind die Steigungen ziemlich langsam zu berechnen (vermutlich, weil ich tf.gather bin mit). Man kann annehmen, dass iA in aufsteigender Reihenfolge sortiert ist (so kann eine bestimmte Zeile A [a,] wiederverwendet werden, bis sich 'a' ändert). Gibt es einen besseren Weg, dies zu tun? Jede Hilfe wird geschätzt.

Bearbeiten: A und B sind tf.Variables hier und ändern für jede Iteration. Precomputing ist keine Option. iA und iB sind jedoch Konstanten und ändern sich nicht.

Bearbeiten: M >> N. Größe von K ist nicht wirklich ein Problem.

Answer 1

Je nachdem, wie dicht Ihre Indizes sind, kann es sinnvoll sein, ein Array von Zählwerten einfach vorab zu berechnen. In numpy würde das

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

Sie werden dann durch diese durch einfache Multiplikation und Addition der indizierten Summe berechnen könnte:

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

Wenn jedoch iA und iB sind beide sehr spärlich, das gibt ganz etwas Leistungs Overhead. Gleiches gilt, wenn iA oder iB nicht festgelegt sind, sondern sich für jede Iteration ändern.