更快替代Tensorflows tf.einsum（）

Question

我需要用三维张量x，y

计算以下张量

tf.einsum("ijk,ljk->ilj",x,y)

不幸的是，这很慢。有没有办法只使用我没想过的matmul操作来重写它？

Answer 1

您可以（隐式）将x广播成iljk的形状，并将y广播成ilkj的形状。 然后可以使用tf.matmul()获得形状iljj，并使用tf.sum（）消除一个j。 这样得到的形状就是ilj。

x = tf.expand_dims(x, axis=1)

y = tf.transpose(y, [0,2,1])
y = tf.expand_dims(y, axis=0)

tf.sum(tf.matmul(x, y), axis=-2)

但是，我不认为这会更快，因为您在此处获得了4D张量。