为何1D场景下tf.boolean_mask比tf.gather+tf.where慢近10倍？

tf.boolean_mask的写法比tf.gather与tf.where的组合更简洁直观，但在一维（1-D）场景下，前者的运行速度明显慢很多，测试代码及结果如下：

import tensorflow as tf

# 定义张量形状
shape = [5000]

# 创建随机布尔掩码m和全1向量v
m = tf.random.uniform(shape) > 0.5
v = tf.ones(shape)

# 使用tf.boolean_mask筛选元素
%timeit tf.boolean_mask(v, m)
# 每循环耗时1.23 ms ± 1.33 µs（7次运行的均值±标准差，每次1000循环）

# 使用tf.gather+tf.where实现相同筛选逻辑
%timeit tf.gather(v, tf.where(m))
# 每循环耗时107 µs ± 349 ns（7次运行的均值±标准差，每次10000循环）

需要注意的是，两种方式的输出结果形状存在细微差异：

tf.boolean_mask(v, m).shape
# TensorShape([2578])
tf.gather(v, tf.where(m)).shape
# TensorShape([2578, 1])

可以通过tf.squeeze去除额外维度修正形状，这会让该组合的运行速度下降约50%：

%timeit tf.squeeze(tf.gather(v, tf.where(m)))
# 每循环耗时149 µs ± 343 ns（7次运行的均值±标准差，每次10000循环）

即便如此，tf.gather+tf.where+tf.squeeze的组合仍比tf.boolean_mask快近10倍。

差异原因解析 #

这两种方式在功能上没有本质差异，核心原因是tf.boolean_mask是通用型算子，它需要适配任意维度张量的掩码筛选逻辑（比如支持在高维张量中指定轴进行掩码操作），没有针对一维场景做特殊优化；而tf.gather+tf.where的组合在一维场景下的执行路径更轻量化：tf.where直接提取掩码为True的索引，tf.gather根据索引直接取值，逻辑完全贴合一维场景的需求，因此性能更优。

另外，张量大小对性能影响显著：当张量形状为[5000000]这种大规模场景时，二者的性能表现基本相当——这是因为大张量下，数据传输和计算的开销成为性能瓶颈，通用实现和专用实现的路径差异被稀释。

内容的提问来源于stack exchange，提问作者fuenfundachtzig