tensorflow - 更快地替代 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 张量。