numpy - `np.dot` 在剩余轴上没有笛卡尔积

Question

根据文档：

对于 N 维，是 的最后一个轴和 的倒数第二个dot轴的和积：ab

dot(a, b)[i,j,k,m] = sum(a[i,j,:] * b[k,:,m])

我想计算最后一个轴a和倒数第二个轴上的和积，b但没有在其余轴上形成笛卡尔积，因为其余轴具有相同的形状。让我用一个例子来说明：

a = np.random.normal(size=(11, 12, 13))
b = np.random.normal(size=(11, 12, 13, 13))
c = np.dot(a, b)
c.shape # = (11, 12, 11, 12, 13)

但我希望形状是(11, 12, 13). 使用广播可以达到预期的效果

c = np.sum(a[..., None] * b, axis=-2)
c.shape # = (11, 12, 13)

但是我的数组相对较大，我想使用并行化 BLAS 实现的强大功能，这些实现似乎np.sum不受np.dot. 关于如何实现这一目标的任何想法？

Answer 1

你可以使用np.einsum-

c = np.einsum('ijk,ijkl->ijl',a,b)

Answer 2

您还可以使用np.matmul：

c = np.matmul(a[..., None, :], b)[..., 0, :]

这相当于Python 3.5+中的 new@运算符：

c = (a[..., None, :] @ b)[..., 0, :]

性能没有太大差异 - 如果np.einsum您的示例数组似乎有什么稍微快一点：

In [1]: %%timeit a = np.random.randn(11, 12, 13); b = np.random.randn(11, 12, 13, 13)
   ....: np.einsum('...i,...ij->...j', a, b)
   ....: 
The slowest run took 5.24 times longer than the fastest. This could mean that an
intermediate result is being cached.
10000 loops, best of 3: 26.7 µs per loop

In [2]: %%timeit a = np.random.randn(11, 12, 13); b = np.random.randn(11, 12, 13, 13)
np.matmul(a[..., None, :], b)[..., 0, :]
   ....: 
10000 loops, best of 3: 28 µs per loop