Python中的快速求和

Question

我想做一系列点积。即

for i in range(N[0]):
  for j in range(N[1]):
    kr[i,j] = dot(k[i,j,:], r[i,j,:])

是否有矢量化方法来执行此操作，例如使用einsum或tensordot？

Answer 1

假设N[0]和N[1]是k和r的前两个维度的长度，

kr = numpy.einsum('...i,...i->...', k, r)

我们指定...启用广播，并沿最后一个轴执行点积。

Answer 2

假设k和r有三个维度，则与以下内容相同：

kr = numpy.sum(k * r, axis=-1)