numpy dot product and matrix product

Question

I'm working with numpy arrays of shape (N,), (N,3) and (N,3,3) which represent sequences of scalars, vectors and matrices in 3D space. I have implemented pointwise dot product, matrix multiplication, and matrix/vector multiplication as follows:

def dot_product(v, w):
    return np.einsum('ij, ij -> i', v, w)

def matrix_vector_product(M, v):
    return np.einsum('ijk, ik -> ij', M, v)

def matrix_matrix_product(A, B):
    return np.einsum('ijk, ikl -> ijl', A, B)

As you can see I use einsum for lack of a better solution. To my surprise I was not able to use np.dot... which seems not suitable for this need. Is there a more numpythonic way to implement these function?

In particular it would be nice if the functions could work also on the shapes (3,) and (3,3) by broadcasting the first missing axis. I think I need ellipsis, but I don't quite understand how to achieve the result.

Answer 1

These operations cannot be reshaped into general BLAS calls and looping BLAS calls would be quite slow for arrays of this size. As such, einsum is likely optimal for this kind of operation.

Your functions can be generalized with ellipses as follows:

def dot_product(v, w):
    return np.einsum('...j,...j->...', v, w)

def matrix_vector_product(M, v):
    return np.einsum('...jk,...k->...j', M, v)

def matrix_matrix_product(A, B):
    return np.einsum('...jk,...kl->...jl', A, B)

Answer 2

Just as working notes, these 3 calculations can also be written as:

np.einsum(A,[0,1,2],B,[0,2,3],[0,1,3])
np.einsum(M,[0,1,2],v,[0,2],[0,1]) 
np.einsum(w,[0,1],v,[0,1],[0])

Or with Ophion's generalization

np.einsum(A,[Ellipsis,1,2], B, ...)

It shouldn't be hard to generate the [0,1,..] lists based on the dimensions of the inputs arrays.

---

By focusing on generalizing the einsum expressions, I missed the fact that what you are trying to reproduce is N small dot products.

np.array([np.dot(i,j) for i,j in zip(a,b)])

It's worth keeping mind that np.dot uses fast compiled code, and focuses on calculations where the arrays are large. Where as your problem is one of calculating many small dot products.

And without extra arguments that define axes, np.dot performs just 2 of the possible combinations, ones which can be expressed as:

np.einsum('i,i', v1, v2)
np.einsum('...ij,...jk->...ik', m1, m2)

An operator version of dot would face the same limitation - no extra parameters to specify how the axes are to be combined.

It may also be instructive to note what tensordot does to generalize dot:

def tensordot(a, b, axes=2):
    ....
    newshape_a = (-1, N2)
    ...
    newshape_b = (N2, -1)
    ....
    at = a.transpose(newaxes_a).reshape(newshape_a)
    bt = b.transpose(newaxes_b).reshape(newshape_b)
    res = dot(at, bt)
    return res.reshape(olda + oldb)

It can perform a dot with summation over several axes. But after the transposing and reshaping is done, the calculation becomes the standard dot with 2d arrays.

---

This could have been flagged as a duplicate issue. People have asking about doing multiple dot products for some time.

Matrix vector multiplication along array axes suggests using numpy.core.umath_tests.matrix_multiply

https://stackoverflow.com/a/24174347/901925 equates:

matrix_multiply(matrices, vectors[..., None])
np.einsum('ijk,ik->ij', matrices, vectors)

The C documentation for matrix_multiply notes:

* This implements the function
* out[k, m, p] = sum_n { in1[k, m, n] * in2[k, n, p] }.

inner1d from the same directory does the same same for (N,n) vectors

inner1d(vector, vector)  
np.einsum('ij,ij->i', vector, vector)
# out[n] = sum_i { in1[n, i] * in2[n, i] }

Both are UFunc, and can handle broadcasting on the right most dimensions. In numpy/core/test/test_ufunc.py these functions are used to exercise the UFunc mechanism.

matrix_multiply(np.ones((4,5,6,2,3)),np.ones((3,2)))

https://stackoverflow.com/a/16704079/901925 adds that this kind of calculation can be done with * and sum, eg

(w*v).sum(-1)
(M*v[...,None]).sum(-1)
(A*B.swapaxes(...)).sum(-1)

---

On further testing, I think inner1d and matrix_multiply match your dot and matrix-matrix product cases, and the matrix-vector case if you add the [...,None]. Looks like they are 2x faster than the einsum versions (on my machine and test arrays).

https://github.com/numpy/numpy/blob/master/doc/neps/return-of-revenge-of-matmul-pep.rst is the discussion of the @ infix operator on numpy. I think the numpy developers are less enthused about this PEP than the Python ones.