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]]

CodePudding user response：

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]])

CodePudding user response：

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)