Tingiskhan Tingiskhan - 2 months ago 14
Python Question

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

Answer

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