how to vectorize this python code to make it more efficient? (in speed)

Question

I am trying to vectorize my code using numpy modules, the original code is like:

m = [M[i,:9].dot(N[:9,i]) for i in xrange(9)]

and I have improved the code as:

m = np.diagonal(M[:9,:9].dot(N[:9,:9]))  

yet this will lead to some unnecessary calculations (especially when the index is much greater than 9). What can I do to improve the efficiency further?

Edit: Basically what I intend to do is to calculate the diagonal elements of the dot product of two matrixes M and N.

Answer 1

You can use np.einsum as we need to keep the first axis of M aligned with the second axis of N, while reducing/losing the leftover axes from the inputs. Thus, we would have an einsum based solution, like so:

m = np.einsum('ij,ji->i',M[:,:9],N[:9,:])