Loopless 3D Array Multiplication

Question

Very similar to https://math.stackexchange.com/q/3615927/419686, but different.

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

Python code:

import numpy as np

np.random.seed(1)
A = np.random.randint(0, 10, size=(5,2,3))
B = np.random.randint(0, 10, size=(6,3,8))

C = np.zeros((5,6,2,8))
for i in range(A.shape[0]):
    for j in range(B.shape[0]):
        C[i,j] = A[i].dot(B[j])

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

@hpaulj Yes. It gets (5,2,6,8). C is different if converted to 1D. — chaosink, Dec 18 '20 at 07:05
That mathstackexchange link does not have very good `numpy` answers. — hpaulj, Dec 18 '20 at 07:15

score 4 · Answer 1 · answered Dec 18 '20 at 07:12

4

In [52]: np.random.seed(1)
    ...: A = np.random.randint(0, 10, size=(5,2,3))
    ...: B = np.random.randint(0, 10, size=(6,3,8))
    ...: 
    ...: C = np.zeros((5,6,2,8))
    ...: for i in range(A.shape[0]):
    ...:     for j in range(B.shape[0]):
    ...:         C[i,j] = A[i].dot(B[j])
    ...:

np.dot does broadcast the outer dimensions:

In [53]: D=np.dot(A,B)
In [54]: C.shape
Out[54]: (5, 6, 2, 8)
In [55]: D.shape
Out[55]: (5, 2, 6, 8)

The axes order is different, but we can easily change that:

In [56]: np.allclose(C, D.transpose(0,2,1,3))
Out[56]: True
In [57]: np.allclose(C, np.swapaxes(D,1,2))
Out[57]: True

From the np.dot docs:

dot(a, b)[i,j,k,m] = sum(a[i,j,:] * b[k,:,m])

answered Dec 18 '20 at 07:12

hpaulj

221,503
14
230
353

Just out of curiosity if there a way to do this using einsum? – Dani Mesejo Dec 18 '20 at 07:26
The einsum expression practically writes itself from the dimensions since they are distinct. – hpaulj Dec 18 '20 at 07:48
Do you have any pointers to good resources for understanding einsum? – Dani Mesejo Dec 18 '20 at 07:49
2

@DaniMesejo `np.einsum('ikl, jlm -> ijkm', A, B)` – Daniel F Dec 18 '20 at 08:41
1

Also, [This answer is literally cited by wikipedia for its clear description of how Einstein summation nototation works](https://stackoverflow.com/questions/26089893/understanding-numpys-einsum/33641428#33641428) – Daniel F Dec 18 '20 at 08:45

score 0 · Answer 2 · answered Dec 18 '20 at 08:42

0

Use np.einsum which is very powerful:

C = np.einsum('aij, bjk -> abik', A, B)

answered Dec 18 '20 at 08:42

chaosink

1,329
13
27

Loopless 3D Array Multiplication

2 Answers2

Linked