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

David Cournapeau cournape at gmail.com
Fri May 22 04:58:15 EDT 2015

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

> Hi,
> 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)
> and
> 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/a155273f/attachment.html>

More information about the NumPy-Discussion mailing list