python - 如何矢量化此 python 代码以使其更高效？ (速度)

Question

我正在尝试使用 numpy 模块对我的代码进行矢量化，原始代码如下:

m = [M[i,:9].dot(N[:9,i]) for i in xrange(9)]

我改进了代码:

m = np.diagonal(M[:9,:9].dot(N[:9,:9]))  

然而这会导致一些不必要的计算(尤其是当索引远大于 9 时)。我可以做些什么来进一步提高效率？

编辑:基本上我打算做的是计算两个矩阵 M 和 N 的点积的对角线元素。

Answer 1

您可以使用 np.einsum因为我们需要保持 M 的第一个轴与 N 的第二个轴对齐，同时减少/丢失输入中的剩余轴。因此，我们将有一个基于 einsum 的解决方案，如下所示:

m = np.einsum('ij,ji->i',M[:,:9],N[:9,:])