块状整体矩阵与单行相乘

Question

我有两个numpy矩阵，A和B：

A = [[0,-1],
     [-2,5],
     [2,6]]
B = [[5,3],
     [2,4]]
Dimension of A: N * K
Dimension of B: K * K

我想做的是与 A的每一行和整个B 进行marix乘法，然后与 A的每一行<< / strong>。然后总结一下。

哪个是

N = A.shape[0]
sum = 0
for i in range(N):
    s = A[i] @ B @ A[i]
    print(s)
    sum += s
return sum

此打印

4
70
224

4 + 70 + 224 = 298是我想要的。 没有for循环，有没有任何优雅或pythonic的方法来实现这一目标？ 谢谢

Answer 1

我们可以使用np.einsum-

np.einsum('ij,jk,ik->',A,B,A)

并使用optimize标志进一步对其进行优化-

np.einsum('ij,jk,ik->',A,B,A,optimize=True)