CODEWITHSUNDEEP
pythonnumpy
chaosink
chaosink

Reputation: 1483

Loopless 2D Array Multiplication

I have 2 matrices (A with shape (2,3) and B with shape (4,3)), and I want to perform some kind of multiplication in order to take a new matrix with shape (2,4,3).

Python code:

import numpy as np

A = np.random.randint(0, 10, size=(2,3))
B = np.random.randint(0, 10, size=(4,3))
C = np.zeros((2,4,3))
for i in range(A.shape[0]):
    for j in range(B.shape[0]):
        C[i,j] = A[i] * B[j]

Is it possible to do the above operation without using a loop?

Upvotes: 0

Views: 57

Answers (2)

hpaulj
hpaulj

Reputation: 231335

In [59]: A = np.random.randint(0, 10, size=(2,3))
    ...: B = np.random.randint(0, 10, size=(4,3))
    ...: C = np.zeros((2,4,3))
    ...: for i in range(A.shape[0]):
    ...:     for j in range(B.shape[0]):
    ...:         C[i,j] = A[i] * B[j]
    ...:

As shown in the other answer, we can do the same with einsum. In contrast to your previous question, you aren't doing a sum-of-products on any dimension. So all dimensions from the left carry over to the right.

In [60]: D = np.einsum('ik,jk->ijk',A,B)
In [61]: np.allclose(C,D)
Out[61]: True

But in numpy we can do this just as well with broadcasting and element-wise multiplication:

In [62]: E = A[:,None,:]*B[None,:,:]
In [63]: np.allclose(C,E)
Out[63]: True

Some comparative timings:

In [64]: timeit D = np.einsum('ik,jk->ijk',A,B)
7.12 µs ± 21.9 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)
In [65]: timeit E = A[:,None,:]*B[None,:,:]
5.18 µs ± 76.7 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)
In [66]: %%timeit
    ...: C = np.zeros((2,4,3))
    ...: for i in range(A.shape[0]):
    ...:     for j in range(B.shape[0]):
    ...:         C[i,j] = A[i] * B[j]
    ...: 
    ...: 

27.5 µs ± 68.5 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

Upvotes: 1

Daniel F
Daniel F

Reputation: 14399

Basically, if you're doing a lot of high-dimensional (i.e. >2D) tensor math, you really want to:

  1. get good at Einstein Notation
  2. Use np.einsum

Good old Albert came up with a way to deal with this for a reason, he literally wrote the book on math in high-dimensional state spaces.

In this case, your answer will be:

C = np.einsum('ik, jk -> ijk', A, B)

Upvotes: 1

Related Questions