利用numpy/pytorch广播计算矩阵积的迹

Question

设A是（nxm）-矩阵，M an（mxm）-矩阵。为矩阵的迹写tr（），我需要计算tr（AM（a^T））。但是，最后的跟踪操作会丢弃大部分计算。我能用纽比或派托克的广播规则来计算AM（A^T）的必要对角线吗？在

更新： 以下是我在Pythorch中计算对角线的解决方案：

torch.sum(torch.sum(A.t()[:,None,:]*M[:,:,None],0)*A.t(),0)

Answer 1

你至少要计算两个矩阵乘积中的一个。随后，您可以在这里使用其中一个答案：What is the best way to compute the trace of a matrix product in numpy?