3D matrix multiplication in numpy

Question

I think yesterday I asked the wrong question . I really want to mutiply two matrices 2x2xN A and B , so that

 C[:,:,i] = dot(A[:,:,i], B[:,:,i]) 

For example, if I have a matrix

 A = np.arange(12).reshape(2, 2, 3) 

How can I get C = A x A with the definition described above? Is there a built-in function for this?

---

Also, if I multiply A (shape 2x2xN) by B (shape 2x2x1, instead of N) , I want to get

 C[:,:,i] = dot(A[:,:,i], B[:,:,1]) 

Answer 1

Try using numpy.einsum , it has a little learning curve, but it should give you what you want. Here is an example to get you started.

 import numpy as np A = np.random.random((2, 2, 3)) B = np.random.random((2, 2, 3)) C1 = np.empty((2, 2, 3)) for i in range(3): C1[:, :, i] = np.dot(A[:, :, i], B[:, :, i]) C2 = np.einsum('ijn,jkn->ikn', A, B) np.allclose(C1, C2)