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

Stack Overflow用户

提问于 2022-09-25 01:02:37

回答 1查看 29关注 0票数 1

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

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

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

回答已采纳

发布于 2022-09-25 02:00:32

你可以用

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

另一个选择是

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

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

票数 1

页面原文内容由Stack Overflow提供。腾讯云小微IT领域专用引擎提供翻译支持

原文链接：

https://stackoverflow.com/questions/73841478

复制

相似问题

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