Neil G - 1 year ago 80

Python Question

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`

Answer Source

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