CODEWITHSUNDEEP
pythonnumpyvectorization
Tingiskhan
Tingiskhan

Reputation: 1345

Dot product between 2D and 3D arrays

Assume that I have two arrays V and Q, where V is (i, j, j) and Q is (j, j). I now wish to compute the dot product of Q with each "row" of V and save the result as an (i, j, j) sized matrix. This is easily done using for-loops by simply iterating over i like

import numpy as np

v = np.random.normal(size=(100, 5, 5))
q = np.random.normal(size=(5, 5))
output = np.zeros_like(v)

for i in range(v.shape[0]):
    output[i] = q.dot(v[i])

However, this is way too slow for my needs, and I'm guessing there is a way to vectorize this operation using either einsum or tensordot, but I haven't managed to figure it out. Could someone please point me in the right direction? Thanks

Upvotes: 1

Views: 866

Answers (1)

Divakar
Divakar

Reputation: 221564

You can certainly use np.tensordot, but need to swap axes afterwards, like so -

out = np.tensordot(v,q,axes=(1,1)).swapaxes(1,2)

With np.einsum, it's a bit more straight-forward, like so -

out = np.einsum('ijk,lj->ilk',v,q)

Upvotes: 1

Related Questions