使用 numpy 高效计算多个点积

Question

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.

Answer 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 .