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

Nadav Horesh nadavh at visionsense.com
Fri May 22 05:50:46 EDT 2015

There was an idea on this list to provide a function the run multiple dot on several vectors/matrices. It seems to be a particular implementation of this proposed function.


On 22 May 2015 11:58, David Cournapeau <cournape at gmail.com> wrote:

On Fri, May 22, 2015 at 5:39 PM, Mathieu Blondel <mathieu at mblondel.org<mailto:mathieu at mblondel.org>> wrote:

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.

Does your implementation use BLAS, or is just a a wrapper around einsum ?


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

More information about the NumPy-Discussion mailing list