CODEWITHSUNDEEP
pythonsumvectorization
yakzo
yakzo

Reputation: 347

Fast Summation in Python

I want to do a series of dot products. Namely

for i in range(N[0]):
  for j in range(N[1]):
    kr[i,j] = dot(k[i,j,:], r[i,j,:])

Is there a vectorized way to do this, for example using einsum or tensordot?

Upvotes: 1

Views: 178

Answers (2)

javidcf
javidcf

Reputation: 59731

Assuming k and r have three dimensions, this is the same as:

kr = numpy.sum(k * r, axis=-1)

Upvotes: 1

user2357112
user2357112

Reputation: 282006

Assuming N[0] and N[1] are the lengths of the first two dimensions of k and r,

kr = numpy.einsum('...i,...i->...', k, r)

We specify ... to enable broadcasting, and perform a dot product along the last axis.

Upvotes: 5

Related Questions