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

Question

根据documentation :

For N dimensions dot is a sum product over the last axis of a and the second-to-last of 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.sum 支持。点。关于如何实现这一目标的任何想法？

Answer 1

您可以使用 np.einsum -

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