Python: Multiplying a (4,1,2)-array with a (4,1) to get a (4,1,2) array

Question

As the title says, I basically want the second matrix (let's call it B) to be a list of multiplications for the rows of the first matrix (which we'll call A).

How do I go about doing this? Example:

A = np.array([[[ 1.,  3.]],
     [[ 1.,  4.]],
     [[ 1.,  5.]],
     [[ 1.,  8.]]])

B = np.array([[0],
     [1],
     [3],
     [8]], dtype=np.int64)

And I want the result to be

C = np.array([[[0., 0.,]],
     [[1.,  4.]],
     [[3.,  15.]],
     [[8.,  64.]]])

Answer 1

For broadcasting to work in this case you'll need to give B a new axis:

>>> A * B[:,None,:]
array([[[  0.,   0.]],
       [[  1.,   4.]],
       [[  3.,  15.]],
       [[  8.,  64.]]])

Answer 2

You might find np.einsum() to be useful with n-dimensional arrays. In this case, you can do

np.einsum('i...,i...->i...',A,B)

to get your desired answer.

np.einsum('ijk,il->ijk',A,B)

also works, and is more explicit.

More details on np.einsum() can be found here.