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

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

Answer 2

来自inner1d的{​​{1}}更快：

重现情节的代码：

numpy.core.umath_tests

Answer 3

即使中等数据大小的速度明显变慢，我也会使用

np.diag(A.dot(B.T))

在开发库时，担心以后在生产环境中运行时，或者在编写单元测试后进行优化。

对于大多数使用您的代码的人来说，这比einsum更容易理解，并且也不要求您通过将计算嵌入迷你DSL字符串中来打破一些最佳实践一些函数调用的参数。

我同意计算非对角元素对于大型案例是值得避免的。尽管如此，对我来说真的非常大，而在einsum中用嵌入式字符串表达计算的可怕代价的折衷非常严重。