How most efficiently compute the diagonal of a matrix product
I want to compute the following:

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)


The first method has the problem that it uses a for loop, which can't be very effecient (I need to do this in pytorch on a GPU millions of times)

The second method computes the full matrix product, when I only need the diagonal entries, so that can't be very efficient either.

I'm wondering whether there exist any clever way of doing this?

Use a manually constructed sum-product. You want the sums of the squares of the individual columns:

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


As Divakar suggests, np.einsum will likely offer a less memory-intensive option, since it does not require the temporary array x * x:

y = np.einsum('ij,ij-&gt;j', x, x)

