使用numpy / pytorch广播计算矩阵产品的跟踪

Question

Let A be an (nxm)-matrix and M an (mxm)-matrix. Writing tr() for the trace of a matrix, I need to compute tr(AM(A^T)). However, the final trace operation throws away most of the computation. Can I use numpy's or pytorch's broadcasting rules to compute only the necessary diagonal of AM(A^T)?

Update: Here is my solution to compute the diagonal in PyTorch:

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

Answer 1

You will have to compute at least one of the two matrix products. Subsequently you can use one of the answers here: What is the best way to compute the trace of a matrix product in numpy?