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 object
s.