Numpy - row-wise outer product of two matrices

Question

I have two numpy arrays: A of shape (b, i) and B of shape (b, o). I would like to compute an array R of shape (b, i, o) where every line l of R contains the outer product of the row l of A and the row l of B. So far what i have is:

import numpy as np

A = np.ones((10, 2))
B = np.ones((10, 6))
R = np.asarray([np.outer(a, b) for a, b in zip(A, B)])
assert R.shape == (10, 2, 6)

I think this method is too slow, because of the zip and the final transformation into a numpy array.

Is there a more efficient way to do it ?

Answer 1

That is possible with numpy.matmul, which can do multiplication of "matrix stacks". In this case we want to multiply a stack of column vectors with a stack of row vectors. First bring matrix A to shape (b, i, 1) and B to shape (b, 1, o). Then use matmul to perform b times the outer product:

import numpy as np

i, b, o = 3, 4, 5

A = np.ones((b, i))
B = np.ones((b, o))

print(np.matmul(A[:, :, np.newaxis], B[:, np.newaxis, :]).shape)  # (4, 3, 5)

An alternative could be to use numpy.einsum, which can directly represent your index notation:

np.einsum('bi,bo->bio', A, B)