如何将这段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,:])