Question

我要进行以下计算，我使用随机数组进行演示：

a = np.random.randint(10, size=(100,3))
b = np.random.randint(10, size=(3,2))

result = np.zeros(100)

for i in range(100):
    result[i] = a[i] @ b @ b.T @ a[i].T

为了加快计算速度，我考虑过用爱因斯坦总和删除for循环。因此，我尝试使用相同的向量进行以下操作：

result = np.einsum('ij,jk,jk,ij->i', a, b, b, a)

我将'i'放在einsum的右侧，因为结果向量显示正确的大小。但是，结果略有不同。

我的问题可以用einsum来解决吗？

弗朗兹

Answer 1

在一个einsum中，它应该是-

np.einsum('ij,jl,kl,ik->i',a,b,b,a)

用matrix-multiplication带入np.dot-

np.einsum('ij,jk,ik->i',a,b.dot(b.T),a)

或更多-

np.einsum('ij,ij->i',a.dot(b.dot(b.T)),a)

使用np.matmul/@-operator in Python 3.x，它转换为-

((a@(b@b.T))[:,None,:] @ a[:,:,None])[:,0,0]

使用np.einsum替换for循环

1 个答案: