1

I am trying to compute many dot products efficiently. This post gets really close to what I am trying to do, but I can't quite get it to work. I have a large list of matrices (a), and a list of vectors (b). I want to do a series of dot product operations between them. This is what works now:

import numpy as np
a # shape (15000,4,4)
b # shape (15000,4)

out = np.empty((15000,4))

for i in range(15000):
    out[i] = np.dot(a[i],b[i])

All my attempts to adapt np.tensordot or np.einsum from the linked post have failed to give me what I want. If anyone sees how to do this I would really appreciate it.

cschlick
  • 61
  • 4

1 Answers1

1

Einstein summation works just fine:

>>> a = np.random.randn(100, 4, 4)
>>> b = np.random.randn(100, 4)
>>> foo = np.einsum('ijk,ik->ij', a, b)
>>> bar = np.zeros_like(foo)
>>> for i, (ai, bi) in enumerate(zip(a, b)):
        bar[i] = np.dot(ai, bi)
>>> np.allclose(foo, bar)
True

To explain the summation a bit, note that you're contracting the last axis of b. So you can think about doing each inner product as if by np.einsum('jk,k->j', a[0], b[0]). But we're doing one for each element of a and b, thus the inclusion of the first axis, which is not contracted. Hence, ijk,ik->ij.

bnaecker
  • 6,152
  • 1
  • 20
  • 33
  • Thank you! That works and einsum is two orders of magnitude faster from a quick test. I guess I just need to really carefully read about einsum. Thanks again. – cschlick Jan 23 '20 at 01:59
  • We can 'batch' process the array on the first dimension with a bit of dimension fiddling: `(a@b[:,:,None])[...,0] ` – hpaulj Jan 23 '20 at 02:20