Scalable solution for dot product of two vectors

Question

The dot product of two vectors can be computed via numpy.dot. Now I want to compute the dot product of an array of vectors:

>>> numpy.arange(15).reshape((5, 3))
array([[ 0,  1,  2],
       [ 3,  4,  5],
       [ 6,  7,  8],
       [ 9, 10, 11],
       [12, 13, 14]])

The vectors are row vectors and the output should be a 1d-array containing the results from the dot products:

array([  5,  50, 149, 302, 509])

For the cross product (numpy.cross) this can be easily achieved specifying the axis keyword. However numpy.dot doesn't have such an option and passing it two 2d-arrays will result in the ordinary matrix product. I also had a look at numpy.tensordot but this doesn't seem to do the job either (being an extended matrix product).

I know that I can compute the dot product per element for 2d-arrays via

>>> numpy.einsum('ij, ji -> i', array2d, array2d.T)

However this solution doesn't work for 1d-arrays (i.e. just a single element). I would like to obtain a solution that works for both 1d-arrays (returning a scalar) and arrays of 1d-arrays (aka 2d-arrays) (returning a 1d-array).

Divakar Divakar · Accepted Answer · 2017-02-18T12:50:13

Use np.einsum with ellipsis(...) to account for arrays with variable number of dimensions, like so -

np.einsum('...i,...i ->...', a, a)

Stating the docs on it -

To enable and control broadcasting, use an ellipsis. Default NumPy-style broadcasting is done by adding an ellipsis to the left of each term, like np.einsum('...ii->...i', a). To take the trace along the first and last axes, you can do np.einsum('i...i', a), or to do a matrix-matrix product with the left-most indices instead of rightmost, you can do np.einsum('ij...,jk...->ik...', a, b).

Sample runs on 2D and 1D arrays -

In [88]: a
Out[88]: 
array([[ 0,  1,  2],
       [ 3,  4,  5],
       [ 6,  7,  8],
       [ 9, 10, 11],
       [12, 13, 14]])

In [89]: np.einsum('...i,...i ->...', a, a) # On 2D array
Out[89]: array([  5,  50, 149, 302, 509])

In [90]: b = a[:,0]

In [91]: b
Out[91]: array([ 0,  3,  6,  9, 12])

In [92]: np.einsum('...i,...i ->...', b,b) # On 1D array
Out[92]: 270

Runtime test -

Since, we need to keep one axis aligned, at least with 2D arrays, one of np.einsum or np.matmul or the latest @ operator would be efficient.

In [95]: a = np.random.rand(1000,1000)

# @unutbu's soln
In [96]: %timeit (a*a).sum(axis=-1)
100 loops, best of 3: 3.63 ms per loop

In [97]: %timeit np.einsum('...i,...i ->...', a, a)
1000 loops, best of 3: 944 µs per loop

In [98]: a = np.random.rand(1000)

# @unutbu's soln
In [99]: %timeit (a*a).sum(axis=-1)
100000 loops, best of 3: 9.11 µs per loop

In [100]: %timeit np.einsum('...i,...i ->...', a, a)
100000 loops, best of 3: 5.59 µs per loop

Scalable solution for dot product of two vectors

1 Answers