两个矩阵的逐行点积的优美表达 [英] Elegant expression for row-wise dot product of two matrices

问题描述

我有两个具有相同尺寸的二维numpy数组A和B，并且正在尝试计算它们的按行点积.我可以做到:

I have two 2-d numpy arrays with the same dimensions, A and B, and am trying to calculate the row-wise dot product of them. I could do:

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

是否还有另一种方法可以使numpy在一步而不是两步中进行逐行点积运算?也许与tensordot?

Is there another way to do this so that numpy is doing the row-wise dot product in one step rather than two? Maybe with tensordot?

推荐答案

这是 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

这篇关于两个矩阵的逐行点积的优美表达的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持IT屋！