python - Scale rows of 3D-tensor

Question

I have an n-by-3-by-3 numpy array A and an n-by-3 numpy array B. I'd now like to multiply every row of every one of the n 3-by-3 matrices with the corresponding scalar in B, i.e.,

import numpy as np

A = np.random.rand(10, 3, 3)
B = np.random.rand(10, 3)

for a, b in zip(A, B):
    a = (a.T * b).T
    print(a)

Can this be done without the loop as well?

Answer 1

You can use NumPy broadcasting to let the elementwise multiplication happen in a vectorized manner after extending B to 3D after adding a singleton dimension at the end with np.newaxis or its alias/shorthand None. Thus, the implementation would be A*B[:,:,None] or simply A*B[...,None].