numpy multiply matrices preserve third axis

Question

if I have k many n,m matrices represented by a n,m,k numpy array, how can I multiply that by k many m,j matrices represented by a 'm,j,k' numpy array at the same time giving me a n,j,k ndarray?

In other words, I need to perform k many matrix multiplications of n,m * m,j = n,j. Is it possible to perform them at once?

EDIT: All of the dimensions vary, but are in general large.

Answer 1

The second solution of @Ophion can do without a loop, and it is faster with larger dimension:

In [65]:

#k,n,m,j=2,4,5,6
k,n,m,j=100,101,102,103
A=np.random.random((n,m,k))
B=np.random.random((m,j,k))
In [66]:

%timeit np.rollaxis(np.array(map(np.dot, np.rollaxis(A,2), np.rollaxis(B,2))), 0, 3)
1 loops, best of 3: 313 ms per loop
In [67]:

%timeit np.einsum('nmk,mjk->njk',A,B)
1 loops, best of 3: 793 ms per loop

And slower than enisum when dimension is small:

In [68]:

k,n,m,j=2,4,5,6
#k,n,m,j=100,101,102,103
A=np.random.random((n,m,k))
B=np.random.random((m,j,k))
In [69]:

%timeit np.rollaxis(np.array(map(np.dot, np.rollaxis(A,2), np.rollaxis(B,2))), 0, 3)
10000 loops, best of 3: 73.7 µs per loop
In [70]:

%timeit np.einsum('nmk,mjk->njk',A,B)
100000 loops, best of 3: 13.5 µs per loop

Sure, this is for python 2.x, in 3.x, be aware that the map returns map objects.

Answer 2

This really depends on the size and shape of your arrays. Where n,m, and j are small you can do something like the following:

import numpy as np
>>> a = np.random.rand(5,2,1E6)
>>> b = np.random.rand(2,5,1E6)
>>> out = np.einsum('nmk,mjk->njk',a,b)
>>> out.shape
(5, 5, 1000000)

If n, m, and j are large you might want to take advantage of a BLAS like so:

>>> a= np.random.rand(1E3,1E2,5)
>>> b= np.random.rand(1E2,1E3,5)
>>> out = np.empty((1E3,1E3,5))
>>> for x in range(out.shape[-1]):
...     out[:,:,x] = np.dot(a[:,:,x], b[:,:,x])

Keep in mind that numpy arrays are row-major. You may want to rearrange your data depending on what you are doing.

Answer 3

Using Daniels code Python 3.10 complains TypeError: 'float' object cannot be interpreted as an integer. This can be fixed casting 1E6 to integer with a = np.random.rand(5,2,int(1E6)).