[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>
wrote:

> 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 ?

David
-------------- 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