如何最有效地计算矩阵乘积的对角线

Question

我想计算以下内容：

import numpy as np
n= 3
m = 2
x = np.random.randn(n,m)
#Method 1
y = np.zeros(m)
for i in range(m):
y[i] = x[:,i] @ x[:,i]
#Method 2
y2 = np.diag(x.T @ x)

第一种方法有一个问题，它使用了for循环，这不是很有效(我需要在GPU上的pytorch中这样做数百万次(

第二种方法计算全矩阵乘积，而我只需要对角线条目，所以这也不是很有效。

我想知道是否有什么聪明的方法可以做到这一点？

Answer 1

使用手动构建的和积。您需要各个列的平方和：

y = (x * x).sum(axis=0)

正如Divakar所建议的，np.einsum可能会提供一个内存占用较少的选项，因为它不需要临时数组x * x:

y = np.einsum('ij,ij->j', x, x)