CODEWITHSUNDEEP
pythonnumpy
chaosink
chaosink

Reputation: 1483

Loopless 3D Array Multiplication

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?

Upvotes: 0

Views: 70

Answers (2)

chaosink
chaosink

Reputation: 1483

Use np.einsum which is very powerful:

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

Upvotes: 0

hpaulj
hpaulj

Reputation: 231335

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])

Upvotes: 4

Related Questions