<div dir="ltr"><div><div><div><div><div><div><div><div>Hi,<br><br></div>I often need to compute the equivalent of<br><br></div>np.diag(np.dot(A, B)).<br><br></div>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<br><br>np.sum(A * B.T, axis=1)<br><br>and<br><br>np.einsum("ij,ji->i", A, B).<br><br></div>The first can allocate quite a lot of temporary memory.<br></div>The second can be quite cryptic for someone not familiar with einsum.<br><div>I assume that einsum does not compute np.dot(A, B), but I haven't verified.<br></div></div><div><br></div>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.<br><br></div>Best,<br></div>Mathieu<br></div>