Question

给出两个矩阵A和B（维度为N×K）和一个索引列表iA =（a1，a2 ... aM）对于A和iB =（b1，b2 ... bM）对于B，I需要做以下事情：

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

Tensorflow中的

索引列表有重复，所以我们不止一次绘制同一行。

我目前的实现如下：

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

然而，渐变的计算速度很慢（可能是因为我使用的是tf.gather）。可以假设iA按升序排序（因此特定行A [a，]可以重复使用直到'a'更改）。有一个更好的方法吗？任何帮助表示赞赏。

编辑：A和B是tf.Variables，并且每次迭代都会更改。预计算不是一种选择。然而，iA和iB是常数，不会改变。

编辑：M＆gt;＆gt; N. K的大小不是真正的问题。

Answer 1

根据索引的密集程度，简单地预先计算一组计数可能是个好主意。在numpy那将是

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

然后，你可以通过简单地乘以这个和求和来计算索引和：

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

但是，如果iA和iB都非常稀疏，则会产生相当大的性能开销。如果iA或iB未修复，则相同，但每次迭代都会更改。