如何用广播计算点积？

Question

a=np.arange(18).reshape(2,3,3)
b=np.arange(6).reshape(2,3)

我想计算点积

a[0]@b[0]
array([ 5, 14, 23])

a[1]@b[1]
array([122, 158, 194])

广播，我试过了

c=a@b[...,None]
c
array([[[  5],
        [ 14],
        [ 23]],

       [[122],
        [158],
        [194]]])

但形状不是我想要的

c.shape
(2, 3, 1)

1. 除了换轴的方式外，如何在计算中得到(2, 3)而不是(2, 3, 1)的形状？

2. 对于广播，为什么[:,None] 不起作用？ ... 在这里是什么意思？

Answer 1

在这种情况下，b[...,None] 与 b[:,:,None] 相同，是一个 (2,3,1) 数组。 ... 的意思是“需要多少:”。

所以点积和是a 的最后一个3 和b 的中间3（倒数第二个）。

您可以使用squeeze 摆脱size 1 维度。

但是对于 (2,3,3) 和 (2,3)，您想要哪个 dot 产品？在einsum 符号中，我可以看到这样做

'ijk,ij->ik'
'ijk,ik->ij'
'ijk,mj->imk'
etc

dot 具有 2 个二维数组的产品定义明确。但是当一个是 3d 时，就会有一些歧义。

In [2]: a=np.arange(18).reshape(2,3,3)
   ...: b=np.arange(6).reshape(2,3)
   ...: 

In [3]: np.einsum('ijk,ik->ij',a,b)
Out[3]: 
array([[  5,  14,  23],
       [122, 158, 194]])

In [4]: np.dot(a,b)
ValueError: shapes (2,3,3) and (2,3) not aligned: 3 (dim 2) != 2 (dim 0)

In [6]: np.dot(a,b[:,:,None]).shape  # 'ijk,kml->ijml'
Out[6]: (2, 3, 2, 1)

In [7]: np.matmul(a,b[:,:,None]).shape   # @ 
Out[7]: (2, 3, 1)
In [8]: np.einsum('ijk,ikm->ijm',a,b[...,None])
Out[8]: 
array([[[  5],
        [ 14],
        [ 23]],

       [[122],
        [158],
        [194]]])

In [12]: np.squeeze(_)   # removing that added dimension
Out[12]: 
array([[  5,  14,  23],
       [122, 158, 194]])

matmul 文档中的相关说明是：

如果任一参数是 N-D，N > 2，则将其视为堆栈 驻留在最后两个索引中的矩阵并相应地广播。

ValueError - 如果a 的最后一个维度与 b 的倒数第二个维度。

---

matmul 中的广播示例是：

In [15]: a@b.T
Out[15]: 
array([[[  5,  14],
        [ 14,  50],
        [ 23,  86]],

       [[ 32, 122],
        [ 41, 158],
        [ 50, 194]]])

In [16]: _.shape
Out[16]: (2, 3, 2)

In [17]: a@b.T[None,:,:]
Out[17]: 
array([[[  5,  14],
        [ 14,  50],
        [ 23,  86]],

       [[ 32, 122],
        [ 41, 158],
        [ 50, 194]]])

更新

我刚刚了解到 optimize=True 现在是 einsum 的默认值，而且这并不总是最快的。

In [108]: %timeit np.einsum('ijk,ik->ij',a,b, optimize=False)
5.66 µs ± 63.2 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

In [109]: %timeit np.einsum('ijk,ik->ij',a,b, optimize=True)
73 µs ± 69.3 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

Einsum optimize fails for basic operation