对于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)
但是我的阵列相对较大,我想使用并行BLAS实现的强大功能,np.sum
似乎不支持但np.dot
支持。关于如何实现这一点的任何想法?
答案 0 :(得分:3)
您可以使用np.einsum
-
c = np.einsum('ijk,ijkl->ijl',a,b)
答案 1 :(得分: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