What is the numpy operation to make a dot product over an axis

Viewed 35

I've got an array (L) of shape (2,2) and an array (W) of shape (2, 5, 3) I'd like to know what is the operation of that does a dot product for each element in axis 2. the result should be of shape (2,5,3). I've tried:

np.malmul(L, W)
ValueError: matmul: Input operand 1 has a mismatch in its core dimension 0

and

np.tensordot(L, W) 
ValueError: shape-mismatch for sum

both return me an error. The slow non pythonic solution is:

W_corr = []
for i in range(W.shape[-1]):
     res_ = L.dot(W[:,:,i])
     W_corr.append(res_)
W_corr = np.moveaxis(np.array(W_corr), 0, -1)

But I'm sure there's a better way. Any idea?

2 Answers

Use .swapaxes() for this:

L = np.random.rand(2, 2)
W = np.random.rand(2, 5, 3)

W_corr = np.dot(W.T, L.T).swapaxes(0, 2)

Swap the first two axes of W, and then do the dot product:

>>> np.random.seed(0)
>>> L = np.random.rand(2, 2)
>>> W = np.random.rand(2, 5, 3)
>>> L.dot(W.transpose(1, 0, 2))   # or W.swapaxes(0, 1)
array([[[0.85473091, 1.05437284, 0.81170348],
        [0.8194622 , 1.0870973 , 0.2950265 ],
        [0.8921741 , 0.39278942, 0.98736769],
        [0.8812003 , 0.33554736, 0.2370251 ],
        [0.56481983, 0.78318688, 0.83360085]],

       [[0.72941859, 0.92255399, 0.69920961],
        [0.78898045, 1.00615784, 0.29557025],
        [0.82590506, 0.39690928, 0.85713066],
        [0.84226213, 0.26876025, 0.19667025],
        [0.43405383, 0.75042139, 0.77877449]]])
Related