Numpy广播机制的条件

前言

书上对广播的判断条件为：“如果两个数组的后缘维度(从末尾开始算起的维度)的轴长度相符或其中一方的长度为1，则认为它们是广播兼容的。”
额，对于书上的，我感觉有点难理解哈。
下面将以尽可能详细的方式说明怎样判断两数组是否可以广播。

判断两个数组是否可广播

1. 写出两数组的shape。
2. 统一shape所含元素的个数。
   • 若两数组的shape所含元素的个数相等（即同维度），则直接执行步骤3；
   • 若两数组shape所含的元素个数不等（即不同维度），在维度较少的shape前面补若干个1，直到两个shape的所含元素个数相等。
3. 在两shape中，推荐从后往前按以下方法逐个维度（即逐个数字）进行比较。每次比较的情况以及处理方法如下：
   • 如果数字大小相同，表示该维度可以广播，继续比较下一个维度。
   • 如果数字大小不同，但两数中至少一个为1，则表示该维度可以广播，继续比较下一个维度。
   • 如果数字大小不同，且两数都不为1，则该维度不可广播，不再继续比较，直接判断两个数组不能广播。
4. 若通过步骤3判断出每个维度都可广播，则这两个数组可以进行广播。

例子

例1

若两数组的shape分别为(2,1) 与 (8,2,1)，判断是否可以广播。`

判断流程

1. 比较维度数量：
   • 第一个数组 (2,1) 有2个维度。
   • 第二个数组 (8,2,1) 有3个维度。
   • 维度数量不同，需要在第一个数组的形状前面补上1，使其维度数量与第二个数组相同。补齐后的形状变为 (1,2,1)。
2. 从最后一个维度开始，向前比较每一对维度：
   • 第三个维度（最右边的维度）：
     • 第一个数组的维度大小是1。
     • 第二个数组的维度大小也是1。
     • 维度大小相同，所以这一对维度可以广播。
   • 第二个维度：
     • 第一个数组的维度大小是2。
     • 第二个数组的维度大小也是2。
     • 维度大小相同，所以这一对维度可以广播。
   • 第一个维度（最左边的维度）：
     • 第一个数组的维度大小是1（补齐后的）。
     • 第二个数组的维度大小是8。
     • 第一个数组该维度大小为1，所以可以广播。
3. 因为每个维度都可广播，所以这俩数组可广播。

验证

a = torch.arange(16).numpy().reshape((8,2,1))
b = torch.arange(2).numpy().reshape((2,1))
a+b
#[[[ 0]
#   [ 2]]
# 
#  [[ 2]
#   [ 4]]
# 
#  [[ 4]
#   [ 6]]
# 
#  [[ 6]
#   [ 8]]
# 
#  [[ 8]
#   [10]]
# 
#  [[10]
#   [12]]
# 
#  [[12]
#   [14]]
# 
#  [[14]
#   [16]]]

例2

若两数组的shape分别为(4,) 与 (3,2,5)，判断是否可以广播。`

判断流程

1. 比较维度数量：
   • 第一个数组 (4,) 有1个维度。
   • 第二个数组 (3,2,5) 有3个维度。
   • 维度数量不同，需要在第一个数组的形状前面补1，使其维度数量与第二个数组相同。补齐后的形状变为 (1，1，4)。
2. 从最后一个维度开始，向前比较每一对维度：
   • 第三个维度（最右边的维度）：
     • 第一个数组的维度大小是4。
     • 第二个数组的维度大小也是5。
     • 两维度大小不同且都不为1，则不可以广播。

验证

a = torch.arange(30).numpy().reshape((3,2,5))
b = torch.arange(4).numpy().reshape((-1))
a+b
#ValueError: operands could not be broadcast together with shapes (3,2,5) (4,)

后话

为什么推荐从后往前

上面一大堆比较，重点无非是大小非1的维度进行比较。前面在统一维度时，是前补1，而大小为1的维度总是可广播的，所以推荐从后往前，尽可能以少的比较次数判断是否可以广播。