python - 调整矩阵数组乘法以使用 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 是不可能的，因为不是 sum-reduced 的轴沿两个单独的轴逐元素相乘。因此，当与 np.tensordot 一起使用时，我们最终会得到两个长度 N 的轴，分别对应于两个输入的第一个轴。