broadcast each row of A along B, avoiding `repeat`

Question

I have 2 2D arrays, with 1 axis of the same dimension:

a = np.array(np.arange(6).reshape((2,3)))
b = np.array(np.arange(12).reshape((3,4)))

I want to multiply and broadcast each row of a with b, that is

b_r = np.repeat(b[:,:,None], 2, axis=2)
ab = a.T[:,None,:] * b_r

Is it possible to do the broadcasting while avoiding the repeat? The idea is to avoid unnecessary memory allocation for the repeat operation.

Answer 1

You can just feed in b[:,:,None] without the repeat, as broadcasting with its very definition would broadcast it for you.

Thus, simply do -

ab = a.T[:,None,:]*b[:,:,None]

We can make it a bit compact though by skipping the trailing : for a and using ... to replace :,: for b, like so -

ab = a.T[:,None]*b[...,None]

---

For the kicks, here's one using np.einsum, which would be little less performant, but more expressive once we get past its string notation -

ab = np.einsum('ij,jk->jki',a,b)