用Numpy计算`np.diagonal(np.dot(A，B)，axis1=1，axis2=2)的有效方法

Question

我有一个numpy数组，A表示shape (n, m, k)，B表示shape (k, m)。我想知道是否有更有效的方法来执行以下操作：

np.diagonal(np.dot(A, B), axis1=1, axis2=2)

由于它正在执行许多在np.dot中不需要的计算(我只需要沿着产生的3-D数组的两个轴对角)。

Answer 1

你可以用

np.einsum('ijk,kj->ij', A, B)

另一个选择是

(A * B.T).sum(axis=-1)

但是在几个不同大小的数组的测试中，einsum版本总是更快。