Search code examples
pythonnumpyarray-broadcasting

Multiplication with two 2D matrix to output 3D matrix

I am wondering any good ways to calculate this type of multiplication.

It's simply multiplying x[i] by x element-wise, and resulting into [2, 2, 3] matrix.

>>> x
array([[0, 1, 2],
       [3, 4, 5]])
>>> output
array([[[ 0,  1,  4],
        [ 0,  4, 10]],

       [[ 0,  4, 10],
        [ 9, 16, 25]]])

I tried with code below and wondering for faster version using numpy.

np.array([
    np.multiply(x[i], x) 
    for i in range(x.shape[0])
])
Solution

  • There are two straightforward ways to do so, the first is using broadcasting, and the second one using einsum. I'd recommed using timeit, to compare the various versions for their speed with the application you have in mind:

    out_broadcast = x[:, None, :] * x
out_einsum = np.einsum('ij,kj->ikj',x,x)