矩阵乘法，但仅在特定的行和列之间

Question

我有三个两个矩阵A和B，我想要的矩阵乘积是diagonal(A.B.A^T)，其中A^T是矩阵的转置。矩阵的维数如下

A - (2^n, n) 
B - (n, n)

n在哪里是任何自然数。

我希望将矩阵A的第一行切片与矩阵B相乘，并将它们的乘积与矩阵A^T的第一列相乘。我不想要A.B.A^T的完整乘积，因为我只想要对角线切片。

在我看来，这可以使用eisum来实现。

关于codereview

的一个相关问题

Answer 1

这里是如何用einsum做到的

np.random.seed(1)

A = np.random.randint(0,10,(8,4))
B = np.random.randint(0,10,(4,4))

# brute force for reference
np.diag(A@B@A.T)
# array([3830,  233, 2835,  958, 3706, 1273, 5478,  934])

# more economical
np.einsum('ij,jk,ik->i',A,B,A)
# array([3830,  233, 2835,  958, 3706, 1273, 5478,  934])