numpy: broadcast multiplication over one common axis of two 2d arrays

Question

I'm looking for a way to element-wise multiply two 2d arrays of shape (a, b) and (b, c), respectively. Over the 'b' axis, which the two arrays have in common.

For instance, an example of what I'd like to broadcast (vectorize) is:

import numpy as np    

# some dummy data
A = np.empty((2, 3))
B = np.empty((3, 4))

# naive implementation
C = np.vstack(np.kron(A[:, i], B[i, :]) for i in [0, 1, 2])

# this should give (3, 2, 4)
C.shape

Does anyone know what to do here? Is there a better way?

Answer 1

With different test cases:

In [56]: A=np.arange(6).reshape((2,3))
In [57]: B=np.arange(12).reshape((3,4))
In [58]: np.vstack([np.kron(A[:,i],B[i,:]) for i in range(3)])
Out[58]: 
array([[ 0,  0,  0,  0,  0,  3,  6,  9],
       [ 4,  5,  6,  7, 16, 20, 24, 28],
       [16, 18, 20, 22, 40, 45, 50, 55]])

A first try with `einsum, preserving all 3 axes (no summing)

In [60]: np.einsum('ij,jk->ijk',A,B)
Out[60]: 
array([[[ 0,  0,  0,  0],
        [ 4,  5,  6,  7],
        [16, 18, 20, 22]],

       [[ 0,  3,  6,  9],
        [16, 20, 24, 28],
        [40, 45, 50, 55]]])

Same numbers, but in a different shape.

I can reorder axes on output, to make a 2x4x3, which can be reshaped to 8,3 and transposed.

In [64]: np.einsum('ij,jk->ikj',A,B).reshape(8,3).T
Out[64]: 
array([[ 0,  0,  0,  0,  0,  3,  6,  9],
       [ 4,  5,  6,  7, 16, 20, 24, 28],
       [16, 18, 20, 22, 40, 45, 50, 55]])

So with another iteration I can get rid of the transpose

In [68]: np.einsum('ij,jk->jik',A,B).reshape(3,8)
Out[68]: 
array([[ 0,  0,  0,  0,  0,  3,  6,  9],
       [ 4,  5,  6,  7, 16, 20, 24, 28],
       [16, 18, 20, 22, 40, 45, 50, 55]])

I should have gotten there right away. A is (2,3), B is (3,4), and I want (3,2,4) reshaped to (3,8). i=2, j=3, k=4 => jik.

So another way of describing the problem,

a_ij * b_jk = c_jik

And since I'm not making use of the sum part of the einsum, regular broadcasted multiplication will also work, with one or more transposes.

Answer 2

credit to @hpaulj for the definitions of A and B
use np.outer and np.stack

A = np.arange(6).reshape((2, 3))
B = np.arange(12).reshape((3, 4))

np.stack([np.outer(A[:, i], B[i, :]) for i in range(A.shape[1])])

[[[ 0  0  0  0]
  [ 0  3  6  9]]

 [[ 4  5  6  7]
  [16 20 24 28]]

 [[16 18 20 22]
  [40 45 50 55]]]

---

and to get the np.einsum in the correct shape

np.einsum('ij, jk->jik', A, B)

[[[ 0  0  0  0]
  [ 0  3  6  9]]

 [[ 4  5  6  7]
  [16 20 24 28]]

 [[16 18 20 22]
  [40 45 50 55]]]

---

Broadcasting and transpose

(A[:, None] * B.T).transpose(2, 0, 1)

[[[ 0  0  0  0]
  [ 0  3  6  9]]

 [[ 4  5  6  7]
  [16 20 24 28]]

 [[16 18 20 22]
  [40 45 50 55]]]

---

shape is (3, 2, 4)

timing