两个矩阵的行方点积的优雅表达

Question

我有两个具有相同尺寸的2-d numpy数组,A和B,我正在尝试计算它们的行方点积.我可以:

np.sum(A * B, axis=1)

还有另一种方法可以做到这一点,以便numpy在一步而不是两步中做行式点积？也许有tensordot？

Answer 1

这是一个很好的应用程序numpy.einsum.

a = np.random.randint(0, 5, size=(6, 4))
b = np.random.randint(0, 5, size=(6, 4))

res1 = np.einsum('ij, ij->i', a, b)
res2 = np.sum(a*b, axis=1)

print(res1)
# [18  6 20  9 16 24]

print(np.allclose(res1, res2))
# True

einsum 也往往会快一点.

a = np.random.normal(size=(5000, 1000))
b = np.random.normal(size=(5000, 1000))

%timeit np.einsum('ij, ij->i', a, b)
# 100 loops, best of 3: 8.4 ms per loop

%timeit np.sum(a*b, axis=1)
# 10 loops, best of 3: 28.4 ms per loop