Multiplying a 3D numpy array by a 2D numpy array

Question

I need to multiply a 3D numpy array by a 2D numpy array.

Let's say the 3D array A has shape (3, 100, 500) and the 2D array B has shape (3, 100) . I need element wise multiplication for each of those 500 axes in the 3D array by the 2D array and then I need to sum along the first axis of the resultant array yielding an array of size (100, 500) .

I can get there with a couple of for loops, but surely there must be a numpy function which will achieve this in 1 line? I have had a look at np.tensordot , np.dot , np.matmul , np.prod and np.sum , but none of these functions will do exactly that.

Answer 1

We can exploit numpy broadcasting:

import numpy as np

a = np.random.rand(3,100,500)
b = np.random.rand(3,100)

# add new axis to b to use numpy broadcasting
b = b[:,:,np.newaxis]
#b.shape = (3,100,1)

# elementwise multiplication
m = a*b
# m.shape = (3,100,500)

# sum over 1st axis
s = np.sum(m, axis=0)

#s.shape = (100,500)

Answer 2

在这种情况下，您可以使用np.einsum轻松表达这些操作：

np.einsum("ijk,ij->jk", A, B)

Answer 3

You can broadcast by adding a new unit axis to the 2D array:

np.sum(A * B[..., None], axis=0)

None in an index introduces a unit dimension at that position, which can be used to align axes for broadcasting. ... is shorthand for : as many times as there are dimensions: in this case it's equivalent to :, : since B is 2D.

An alternative way to write it would be

(A * B.reshape(*B.shape, 1)).sum(axis=0)

Answer 4

You can try the following which should work,

np.sum(A.T*B.T,axis=-1).T

This will give you shape (100,500)