同时计算多个相同维度的点积

Question

如果我们计算 4 元素向量和 4x4 矩阵之间的一个点积，我们会这样做

import numpy as np

x = np.random.rand(4)

y = np.random.rand(4,4)

print(x.dot(y))

但是，如果我想做同样的事情，但同时对 10 个不同的向量和矩阵（但具有相同的维度）怎么办？我试过了：

import numpy as np

x = np.random.rand(10,4)

y = np.random.rand(10,4,4)

print(x.dot(y))

但是没有用。我也尝试了其他方法，但我永远无法获得所需的输出维度 (10,4)。不使用 for 循环的最佳方法是什么？

Answer 1

这似乎有效：

import numpy as np

x = np.random.rand(10,4)

y = np.random.rand(10,4,4)

np.einsum("ij,ijk->ik", x, y)

Answer 2

你可以使用矩阵乘法。

@ 运算符广播除最后两个维度之外的所有内容。我们可以通过临时插入一个额外的维度来使用这种行为。这会将 x 转换为 10 个 1x4 矩阵的堆栈。由于矩阵乘法的结果与上一步的形状相同，我们删除了中间维度。

>>> np.allclose((x[:, None] @ y)[:, 0], np.einsum("ij,ijk->ik", x, y))
# True

此方法或多或少等同于 einsum 方法，但在我的机器上速度稍快：

>>> timeit(lambda:np.einsum("ij,ijk->ik", x, y))
3.4603776139992988
>>> timeit(lambda:(x[:,None]@y)[:,0])
2.62706138700014

Computing multiple dot products of same dimension at same time