Loopless 2D Array Multiplication

Question

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?

What's wrong with your previous question: https://stackoverflow.com/questions/65352886/loopless-3d-array-multiplication ? — Dani Mesejo, Dec 18 '20 at 07:56
Does this answer your question? [Loopless 3D Array Multiplication](https://stackoverflow.com/questions/65352886/loopless-3d-array-multiplication) — Jonti, Dec 18 '20 at 07:58
@MetallimaX, see my answer for why we try to avoid (Python level) loops in `numpy`. With `numpy` we try to perform the loops in compiled code, where they are much faster. That's what we commonly mean by 'vectorization', treating the array(s) as a whole thing, rather than a collection of things. — hpaulj, Dec 18 '20 at 16:30
That make sense and should have been clearly exposed in the question. — MetallimaX, Dec 18 '20 at 21:32

score 1 · Answer 1 · answered Dec 18 '20 at 08:28

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

get good at Einstein Notation
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)

score 1 · Answer 2 · answered Dec 18 '20 at 16:27

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)

Loopless 2D Array Multiplication

2 Answers2