Numpy Matrix Multiplication Broadcast

Question

I have an array comprised of N 3x3 arrays (a collection of matrices, although the data type is np.ndarray) and I have an array comprised of N 3x1 arrays (a collection of vectors). What I want to do is multiply each matrix by each vector, so I expect to get back N 3x1 arrays.

Simple example:

A = np.ones((6,3,3))
B = np.ones((6,3,1))
np.dot(A,B) # This gives me a 6x3x6x1 array, which is not what I want
np.array(map(np.dot,A,B)) # This gives me exactly what I want, but I don't want to have to rely on map

I've tired all kinds of reshaping, explored einsum, etc., but can't get this to work the way I want it to. How do I get this to work with numpy broadcasting? This operation will ultimately need to be done many thousands of times, and I don't want map or list comprehension operations to slow things down.

Answer 1

You can use np.einsum to calculate the dot products and create the matrix of the desired shape:

np.einsum('ijk,ikl->ijl', A, B)