numpy multidimensional (3d) matrix multiplication

Question

I get two 3d matrix A (32x3x3) and B(32x3x3), and I want to get matrix C with dimension 32x3x3. The calculation can be done using loop like:

a = numpy.random.rand(32, 3, 3)
b = numpy.random.rand(32, 3, 3)
c = numpy.random.rand(32, 3, 3)

for i in range(32):
    c[i] = numpy.dot(a[i], b[i])

I believe there must be a more efficient one-line solution to this problem. Can anybody help, thanks.

Answer 1

You could do this using np.einsum:

In [142]: old = orig(a,b)

In [143]: new = np.einsum('ijk,ikl->ijl', a, b)

In [144]: np.allclose(old, new)
Out[144]: True

One advantage of using einsum is that you can almost read off what it's doing from the indices: leave the first axis alone (i), and perform a matrix multiplication on the last two (jk,kl->jl)).