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

Question

我有三个两个矩阵A，B我想要的矩阵乘积是diagonal(A.B.A^T)，矩阵的转置在哪里A^T。矩阵的维度如下

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

其中 isn是任意自然数。

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

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

关于codereview的一个相关问题

标签： pythonnumpynumpy-einsum

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])