调整矩阵数组乘法以使用 Numpy Tensordot

Question

我正在尝试加快我的代码以执行一些数值计算，其中我需要将 3 个矩阵与一个数组相乘。问题的结构如下：

• 数组为 (N, 10) 的形状
• 第一个矩阵沿数组的动态维度保持不变，形状为 (10, 10)
• 其他两个矩阵沿数组的第一个维度变化，具有 (N, 10, 10) 形状
• 计算的结果应该是(N, shape)的数组

我已经使用有效的 for 循环实现了一个解决方案，但我希望有更好的性能，所以我尝试使用 numpy 函数。我试过使用 numpy.tensordot 但是当将动态矩阵与数组相乘时，我得到的形状是 (N, 10, N) 而不是 (N, 10)

我的 for 循环如下：

res = np.zeros(temp_rho.shape, dtype=np.complex128)
for i in range(temp_rho.shape[0]):
    res[i] = np.dot(self.constMatrix, temp_rho[i])  
    res[i] += np.dot(self.dinMat1[i], temp_rho[i])
    res[i] += np.dot(self.dinMat2[i], np.conj(temp_rho[i]))
#temp_rho.shape = (N, 10)
#res.shape = (N, 10)
#self.constMatrix.shape = (10, 10)
#self.dinMat1.shape = (N, 10, 10)
#self.dinMat2.shape = (N, 10, 10)

这段代码应该如何实现numpy的点积，返回正确的维度？

Answer 1

这是一种结合使用np.dot 和np.einsum 的方法-

parte1 = constMatrix.dot(temp_rho.T).T
parte2 = np.einsum('ijk,ik->ij',dinMat1, temp_rho)
parte3 = np.einsum('ijk,ik->ij',dinMat2, np.conj(temp_rho))
out = parte1 + parte2 + parte3

获取parte1 的另一种方法是使用np.tensordot -

parte1 = np.tensordot(temp_rho, constMatrix, axes=([1],[1]))

为什么numpy.tensordot 不适用于后面的两个减和？

好吧，我们需要保持dinMat1/dinMat2 之间的第一个轴与temp_rho/np.conj(temp_rho) 的第一个轴对齐，这是不可能的，tensordot 作为不是@ 的轴987654336@ 沿两个单独的轴逐元素相乘。因此，当与np.tensordot 一起使用时，我们最终会得到两个长度为N 的轴，对应于两个输入中的第一个轴。