Neil G - 1 year ago 95
Python Question

# Elegant expression for row-wise dot product of two matrices

I have two 2-d numpy arrays with the same dimensions, A and B, and am trying to calculate the row-wise dot product of them. I could do:

``````np.sum(A * B, axis=1)
``````

Is there another way to do this so that numpy is doing the row-wise dot product in one step rather than two? Maybe with
`tensordot`
?

This is a good application for `numpy.einsum`.

``````a = np.random.randint(0, 5, size=(6, 4))
b = np.random.randint(0, 5, size=(6, 4))

res1 = np.einsum('ij, ij->i', a, b)
res2 = np.sum(a*b, axis=1)

print(res1)
# [18  6 20  9 16 24]

print(np.allclose(res1, res2))
# True
``````

`einsum` also tends to be a bit faster.

``````a = np.random.normal(size=(5000, 1000))
b = np.random.normal(size=(5000, 1000))

%timeit np.einsum('ij, ij->i', a, b)
# 100 loops, best of 3: 8.4 ms per loop

%timeit np.sum(a*b, axis=1)
# 10 loops, best of 3: 28.4 ms per loop
``````
Recommended from our users: Dynamic Network Monitoring from WhatsUp Gold from IPSwitch. Free Download