Matrix multiplication between 2d with 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]])