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

我要计算以下内容：

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上的py​​torch中执行此操作数百万次）

当我只需要对角线项时，第二种方法将计算全矩阵乘积，因此也不是非常有效。

我想知道是否存在任何巧妙的方法？