NumPy: 广播机制

2026-09-30 · numpy · zh · raw

上一篇讨论了 ndarray 如何通过形状与步长访问内存。[memory-model] 广播沿用这一机制，使形状不同的数组能够参与同一次逐元素运算，而不必先将输入复制成相同大小。

本文从二维数组与一维数组的相加开始，观察广播后的索引如何对应原始数据，再对照迭代器源码，分析零步长的设置与实际运算中的内存开销。

实验使用 NumPy 2.3.5，源码固定到 v2.3.5。以下只讨论普通数值 ndarray 的逐元素运算。

1. 形状对齐与元素对应

对一个三行四列的数组，分别给四列加上不同的数：

import numpy as np

a = np.arange(12, dtype=np.int32).reshape(3, 4)
b = np.array([10, 20, 30, 40], dtype=np.int32)
c = a + b

print(c)

[[10 21 32 43]
 [14 25 36 47]
 [18 29 40 51]]

a 的形状为 (3, 4)，b 为 (4,)。结果中的每个元素满足：

c_{i, j} = a_{i, j} + b_{j} .

(1)

同一个 b[j] 参与三行的计算，行索引并不影响它的取值。

NumPy 按照从右向左的顺序对齐各轴。两个轴长度相等时保留该长度；其中一个长度为一时，采用另一个轴的长度；缺失的前导轴按长度一处理。其余情况不兼容。[broadcasting]

因此，(4,) 在本次运算中按 (1, 4) 对齐，随后与 (3, 4) 形成共同形状。这个对齐只是运算规则，不会把原数组 b.shape 改成二维。

可以单独检查形状，不执行加法：[broadcast-shapes]

print(np.broadcast_shapes(a.shape, b.shape))  # (3, 4)

右对齐也意味着 NumPy 不会根据“行”或“列”的用途猜测对应关系。一个形状为 (3,) 的数组不能直接与 a 相加，因为尾轴的 3 与 4 不兼容。若要分别给三行加上不同的值，应显式保留列轴：

per_row = np.array([100, 200, 300], dtype=np.int32)

print(per_row[:, None].shape)                         # (3, 1)
print(np.broadcast_shapes(a.shape, per_row[:, None].shape))  # (3, 4)

这不是将三元素数组补成四元素数组，而是声明它的三个值对应哪一个轴。广播也不是循环填充：长度为二的轴，不能仅因四是二的倍数就广播为长度四。

2. 零步长与数据复用

形状对齐规定了哪些元素参与运算，但还没有说明它们如何存储。若真把 b 复制成三行，列步长应当仍为四字节，行步长则是十六字节。

用 np.broadcast_to 显式构造相应的广播视图：[broadcast-to]

bb = np.broadcast_to(b, a.shape)

print(bb.shape)                        # (3, 4)
print(bb.strides)                      # (0, 4)
print(bb.ctypes.data == b.ctypes.data)  # True
assert np.shares_memory(b, bb)

行步长不是十六，而是零。代入上一篇的寻址公式：

{address}_{b b} (i, j) = {data}_{b} + i \times 0 + j \times 4 = {data}_{b} + 4 j .

(2)

行索引变化时，地址不变。所谓三行数据，实际上是三个逻辑行访问同一份四元素存储。

以第三列为例，三行的字节偏移均为八。仍可以用裸地址读取验证：

import ctypes

offsets = [i * bb.strides[0] + 2 * bb.strides[1] for i in range(3)]
print(offsets)  # [8, 8, 8]

address = bb.ctypes.data + offsets[2]
value = ctypes.c_int32.from_address(address).value
assert value == bb[2, 2] == 30

这里保留了原数组，只读取已知有效的位置。ctypes 不检查数组边界，地址整数本身也不会维持数组的生命周期。[ctypes]

bb 的形状与步长属于这个新视图，原来的 b 仍是一维数组，步长仍为 (4,)。广播没有修改原数组的布局，也没有复制它的元素；新视图及其元数据仍然需要内存。

逻辑大小与存储大小

此时查看 nbytes，会得到两个不同的数：

print(b.nbytes)   # 16
print(bb.nbytes)  # 48

nbytes 按逻辑元素数乘以元素大小计算，不统计不同索引实际占用了多少个互异的内存位置。[nbytes]

bb 有十二个逻辑元素，所以报告四十八字节；这些索引却只访问 b 中的四个整数。因此，不能用 bb.nbytes - b.nbytes 判断广播额外分配了多少内存。

3. 两侧广播的地址映射

广播不必由一个较小数组单方面匹配另一个较大数组。两个输入可以分别沿不同的轴复用数据，共同形成比两者都大的结果。

继续使用 b，再引入一个三行一列的数组：

r = np.arange(3, dtype=np.int32).reshape(3, 1)
rr = np.broadcast_to(r, (3, 4))

print(rr.strides)  # (4, 0)
print(bb.strides)  # (0, 4)
print(r + b)

[[10 20 30 40]
 [11 21 31 41]
 [12 22 32 42]]

r 沿列方向复用，所以列步长为零；b 沿行方向复用，所以行步长为零。对于结果位置 (i, j)，两个输入地址分别为：

\begin{array}{l} {address}_{r r} (i, j) & = {data}_{r} + 4 i, \\ {address}_{b b} (i, j) & = {data}_{b} + 4 j . \end{array}

(3)

两个操作数使用同一组逻辑索引，但各自决定哪些索引影响地址。结果 (2, 1) 从 r 的偏移八字节处读取 2，从 b 的偏移四字节处读取 20，相加得到 22。

长度为一的轴只有索引零。沿这个轴广播，意味着结果坐标无论取何值，输入坐标都保持为零；在地址计算中，这恰好可以由零步长表示。缺失的前导轴采用同样的处理。

因此，广播后的形状描述共同的迭代范围，各操作数的有效步长则描述结果索引如何映射到自身数据。右对齐是 NumPy 规定的对应规则，零步长是实现数据复用的一种直接表示。[ufunc-basics]

rr 和 bb 的三行四列逻辑视图，以及分别只有三个和四个整数的实际存储；索引 (2, 1) 读取偏移 8 和 4 字节的元素，相加得到 22。 — 图 F01：两个广播视图形状相同，但零步长轴不同。同一逻辑索引分别映射到各自输入的存储位置，无须复制成两个三行四列的缓冲区。

4. 广播视图的只读约束

零步长使同一数组内部的多个逻辑位置指向同一地址。在 bb 中，(0, 0)、(1, 0) 和 (2, 0) 并不是三个独立的存储位置。

如果允许像普通独立数组那样赋值，修改一个位置便会同时改变其他位置的读取结果。np.broadcast_to 因此返回只读视图，尝试写入会抛出异常。[broadcast-to]

print(bb.flags.writeable)  # False

try:
    bb[0, 0] = 99
except ValueError as exc:
    print(exc)  # assignment destination is read-only

只读约束作用于这个视图，不会冻结底层数据。原数组 b 仍然可写，其修改会反映到所有共享该位置的逻辑元素：

b[0] = 99
print(bb[:, 0])  # [99 99 99]
b[0] = 10       # 恢复前面的实验数据

这里没有发生三次赋值，只改动了 b 中的一个整数。若需要一个可独立修改的 (3, 4) 数组，应使用 bb.copy()，而不是尝试强行开启广播视图的写权限。

5. 迭代器中的步长设置

前面的实验可以在源码中找到对应。NumPy v2.3.5 的 numpy/lib/_stride_tricks_impl.py 中，broadcast_to 调用内部函数 _broadcast_to，后者通过 np.nditer 创建迭代器，再取出其数组视图。[broadcast-source]

以下摘录构造与取视图的部分；此处 extras 是空列表：

it = np.nditer(
    (array,), flags=['multi_index', 'refs_ok', 'zerosize_ok'] + extras,
    op_flags=['readonly'], itershape=shape, order='C')
with it:
    broadcast = it.itviews[0]

itershape 指定目标迭代形状，op_flags 指定输入只读，multi_index 要求保留多维索引。itviews 提供与迭代器访问布局对应的数组视图。[nditer]

这里没有遍历目标形状的所有元素，也没有逐行复制。构造迭代器时已经建立了访问描述，返回视图即可暴露这套描述。

继续进入 numpy/_core/src/multiarray/nditer_constr.c。npyiter_fill_axisdata 先确定广播形状，再填写各轴上每个操作数的步长。在未指定自定义轴映射的分支中，包含以下判断，省略分支内的后续检查：[iterator-source]

else if (idim >= ondim ||
                PyArray_DIM(op_cur, ondim-idim-1) == 1) {
    strides[iop] = 0;
    /* 后续检查省略。 */
}

idim 从最右侧的轴开始计数，ondim 是当前操作数的维数。条件分别对应两种情况：当前操作数没有这个前导轴，或者对应轴长度为一。两者都使用零步长；对应的非广播分支则读取原数组的轴步长。

需要注意，strides[iop] 是迭代器当前轴上，第 iop 个操作数的有效步长，不是直接改写该数组的 strides 属性。一次迭代可以为多个输入分别保存访问方式，而不改变任何原数组的形状与步长。

这也说明，np.broadcast_to 只是将广播访问方式显式呈现为一个视图；不能据此认为 a + b 必须先调用两次 broadcast_to。逐元素加法由 ufunc 执行，其迭代机制同样能够沿广播轴保持输入地址不变。[ufunc-basics]

6. 结果存储与计算开销

输入不需要扩张，并不意味着加法结果也只是视图。回到最初的 c = a + b：

print(c.nbytes)  # 48
assert not np.shares_memory(c, a)
assert not np.shares_memory(c, b)

结果中的十二个整数需要独立保存。对于这里的普通数组，np.add 未提供 out 时会分配结果数组；提供 out 则写入指定的输出存储。[add]

out = np.empty_like(a)
result = np.add(a, b, out=out)

assert result is out
assert np.array_equal(out, c)

out 允许复用事先分配的结果空间，不会使十二个结果只占四个元素的位置。也不能将它理解为“整个运算保证不使用临时内存”：类型转换、对齐或输入输出重叠等情况，仍可能需要缓冲或临时副本。[ufunc-basics] [ufunc-reference]

更明显的例子是两个一万元素的 float64 向量。将它们分别表示为 (10000, 1) 和 (1, 10000)，两份输入数据合计只有 160,000 字节；但相加所得形状是 (10000, 10000)，完整结果需要：

10000 \times 10000 \times 8 = 800, 000, 000 bytes .

(4)

即十进制的 800 MB，尚未计入对象和其他开销。这个大小可以只通过形状与 dtype 计算，不必真的分配数组。

广播省去的是为重复使用输入而制造的副本，不会省去结果的元素数。只要要求完整保留这次加法的稠密结果，就仍需与结果规模成正比的存储和写入工作。较大的广播表达式因此可能产生较大的结果或中间数组，即便每份输入本身都很小。[broadcasting]

结语

广播规定不同形状的数组如何对应元素；零步长使某些结果坐标不再影响输入地址，从而复用原始数据。输入的存储规模与输出的逻辑规模，应分别考虑。

广播与向量化也不是同一概念。前者解决元素对应，后者通常指通过数组级操作替代 Python 中的逐元素循环。同形数组的相加同样是向量化操作，却不需要扩张任何轴。至于底层循环是否使用 SIMD，则属于执行实现，不能仅凭发生了广播作出判断。[ufunc-basics] [simd]

参考资料

[memory-model] NumPy: ndarray 的内存模型

[broadcasting] NumPy：Broadcasting

[broadcast-shapes] numpy.broadcast_shapes

[broadcast-to] numpy.broadcast_to

[ctypes] numpy.ndarray.ctypes

[nbytes] numpy.ndarray.nbytes

[broadcast-source] NumPy v2.3.5：_broadcast_to

[nditer] numpy.nditer

[iterator-source] NumPy v2.3.5：npyiter_fill_axisdata

[add] numpy.add

[ufunc-basics] NumPy：Universal functions basics

[ufunc-reference] NumPy：Universal functions

[simd] NumPy：CPU/SIMD optimizations

[end]