Numpy Array Broadcasting具有不同的尺寸

Question

我对numpy的广播规则感到有些困惑.假设您想要执行更高维数组的轴方向标量积,以将数组维度减少一(基本上沿一个轴执行加权求和):

from numpy import *

A = ones((3,3,2))
v = array([1,2])

B = zeros((3,3))

# V01: this works
B[0,0] = v.dot(A[0,0])

# V02: this works
B[:,:] = v[0]*A[:,:,0] + v[1]*A[:,:,1] 

# V03: this doesn't
B[:,:] = v.dot(A[:,:]) 

为什么V03不起作用？

干杯

Answer 1

np.dot(a, b)在 a 的最后一个轴和 b 的倒数第二个轴上运行。因此，对于您问题中的特定情况，您始终可以选择：

>>> a.dot(v)
array([[ 3.,  3.,  3.],
       [ 3.,  3.,  3.],
       [ 3.,  3.,  3.]])

如果您想保持v.dot(a)顺序，则需要将轴放置到位，这可以通过以下方式轻松实现np.rollaxis：

>>> v.dot(np.rollaxis(a, 2, 1))
array([[ 3.,  3.,  3.],
       [ 3.,  3.,  3.],
       [ 3.,  3.,  3.]])

我不太喜欢np.dot，除非是为了明显的矩阵或向量乘法，因为它在使用可选参数时对输出数据类型非常严格out。Joe Kington 已经提到过，但是如果您打算做此类事情，请习惯np.einsum：一旦您掌握了语法的窍门，它将减少您花在担心重塑事物上的时间到最少:

>>> a = np.ones((3, 3, 2))
>>> np.einsum('i, jki', v, a)
array([[ 3.,  3.,  3.],
       [ 3.,  3.,  3.],
       [ 3.,  3.,  3.]])

并不是说它在这种情况下太相关，但它也快得离谱：

In [4]: %timeit a.dot(v)
100000 loops, best of 3: 2.43 us per loop

In [5]: %timeit v.dot(np.rollaxis(a, 2, 1))
100000 loops, best of 3: 4.49 us per loop

In [7]: %timeit np.tensordot(v, a, axes=(0, 2))
100000 loops, best of 3: 14.9 us per loop

In [8]: %timeit np.einsum('i, jki', v, a)
100000 loops, best of 3: 2.91 us per loop