将 3D 矩阵与 2D 矩阵相乘，返回与 2D 矩阵具有相同维度的矩阵

Question

In Numpy, say I have a matrix A of dimensions ixjxk, and a matrix B of dimensions kx i. How do I get back a product with dimensions kxi 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!

Answer 1

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.

Answer 2

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)