2d 与 3d 之间的矩阵乘法？

Question

Here is the two matrixs:
a's shape is (2, 2) and b's shape is (2, 2, 3)
I want to get c whose shape is (2, 3)

How can I get c using a and b?

a = array([[0.5, 0.5],
           [0.6, 0.4]])

b = array([[[1, 2, 1],
            [1, 3, 1]],

           [[2, 1, 2],
            [3, 1, 3]]])

c = array([[1. , 2.5, 1. ],
           [2.4 , 1.2, 2.4 ]])

# c = [[0.5*1+0.5*1, 0.5*2+0.5*3, 0.5*1+0.5*1],
       [0.6*2+0.4*3, 0.6*1+0.4*1, 0.6*2+0.4*3]]

# [0.5, 0.5] * [[1, 2, 1],
                [1, 3, 1]]
# [0.6, 0.4] * [[2, 1, 2],
                [3, 1, 3]]

Answer 1

Using einsum

Try np.einsum ( documentation ). If you want to know more about how np.einsum works, then check this old answer of mine which breaks down how its working -

np.einsum('ij,ijk->ik',a,b)

array([[1. , 2.5, 1. ],
       [2.4, 1. , 2.4]])

Using broadcasting

The einsum above is equivalent to the following multiply->reduce->transpose

Note: a[:,:,None] adds an additional axis to matrix a such that (2,2) -> (2,2,1). This allows it to be broadcasted in operations with b which is of the shape (2,2,3).

(a[:,:,None]*b).sum(1)

array([[1. , 2.5, 1. ],
       [2.4, 1. , 2.4]])

Using Tensordot

Check out tensordot documentation here

np.tensordot(a,b, axes=[1,1]).diagonal().T

array([[1. , 2.5, 1. ],
       [2.4, 1. , 2.4]])

Answer 2

The relatively new matmul is designed to handle 'batch' operations like this. The first of 3 dimensions is the batch dimension, so we have to adjust a to be 3d.

In [156]: a = np.array([[0.5, 0.5],
     ...:            [0.6, 0.4]])
     ...: 
     ...: b = np.array([[[1, 2, 1],
     ...:             [1, 3, 1]],
     ...: 
     ...:            [[2, 1, 2],
     ...:             [3, 1, 3]]])
In [157]: (a[:,None]@b)[:,0]
Out[157]: 
array([[1. , 2.5, 1. ],
       [2.4, 1. , 2.4]])

In einsum terms this is

np.einsum('ilj,ijk->ik',a[:,None],b)

with the added l dimension (which is later removed from the result)