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

Matrix multiplication but only between specific rows and columns

我有三个二元矩阵AB,我要的矩阵乘积是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

的相关问题

这是使用 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])