Multiply a 3D matrix with a 2D matrix, return a matrix with the same dimensions as the 2D one

user_11485718 0
Viewed 51

In Numpy, say I have a matrix A of dimensions i x j x k, and a matrix B of dimensions k x i. How do I get back a product with dimensions k x i without using a loop?
For example:

 Let         A = [[[1 2],
                   [3 4]],
                  [[5 6],
                   [7 8]]]

             B = [[a b],
                  [c d]]

I would like to get:

             C = [[a+2c  5b+6d],
                  [3a+4c 7b+8d]]

My current solution is

np.diagonal(np.dot(A, B), axis1=0, axis2=2)

However, the problem is I'm working with a large data set (A and B have huge dimensions) so np.dot(A, B) would lead to MemoryError. Therefore I want to figure out a better way to solve this without computing the dot product.

I have looked into functions like einsum and tensordot but I didn't find what I need (or maybe I missed something). I would appreciate if someone could help me out. Thank you!

vectorization matrix numpy python
2 Answers

You can do this by analyzing the input and output dimensions.

A.shape -> i, j, k
B.shape -> k, i

Your example is not very good, but if you look carefully at what you're asking the output to be, the shape of B must match the first and last dimensions of A.

The sum reduction is happening along the last axis of A and the first axis of B:

np.einsum('ijk,ki->ji', A, B)

einsum can be intimidating, but doing this sort of analysis can save you a lot of the frustration caused by trial and error.

user_29887300

Looking at your output, C[:,0] is

A[0] @ B[0].T

and similarly, C[:,1] is A[1] @ B[1].T.

With that in mind, the np.einsum formula is:

C = np.einsum('ijk,ik->ji', A,B)
user_42384080
Related