`np.dot` 在剩余轴上没有笛卡尔积答案

【问题标题】：`np.dot` without cartesian product on remaining axes`np.dot` 在剩余轴上没有笛卡尔积
【发布时间】：2016-05-05 23:54:51
【问题描述】：

根据documentation：

对于 N 个维度，dot 是 a 的最后一个轴和 b 的倒数第二个轴的和积：
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)

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

【问题讨论】：

标签： numpy array-broadcasting

【解决方案1】：

你可以使用np.einsum -

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

【讨论】：

我真的很喜欢 einsum，但它不能使用 BLAS 进行并行化。
@TillHoffmann 好吧，AFAIK，像 np.dot 和 np.tensordot 这样使用 blas 的基于点的函数不允许对齐轴，比如 einsum 在这里我们保留前两个轴在 a 和 b 之间对齐。在性能方面，对于样本数据大小，这种einsum 方法被证明比4x+ 更好:)
太好了，感谢您的澄清。对于任何感兴趣的人，np.einsum('...i,...ij->...j', a, b) 与维度无关。

【解决方案2】：

你也可以使用np.matmul:

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

这相当于 Python 3.5+ 中的the new @ operator：

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

【讨论】：