[Numpy-discussion] np.diag(np.dot(A, B))

Mathieu Blondel mathieu at mblondel.org
Fri May 22 04:39:00 EDT 2015


I often need to compute the equivalent of

np.diag(np.dot(A, B)).

Computing np.dot(A, B) is highly inefficient if you only need the diagonal
entries. Two more efficient ways of computing the same thing are

np.sum(A * B.T, axis=1)


np.einsum("ij,ji->i", A, B).

The first can allocate quite a lot of temporary memory.
The second can be quite cryptic for someone not familiar with einsum.
I assume that einsum does not compute np.dot(A, B), but I haven't verified.

Since this is is quite a recurrent pattern, I was wondering if it would be
worth adding a dedicated function to NumPy and SciPy's sparse module. A
possible name would be "diagdot". The best performance would be obtained
when A is C-style and B fortran-style.

-------------- next part --------------
An HTML attachment was scrubbed...
URL: <http://mail.python.org/pipermail/numpy-discussion/attachments/20150522/e1d49611/attachment.html>

More information about the NumPy-Discussion mailing list