TensorFlow中使用tf.boolean_mask如何保持原张量维度？

我完全懂你的困扰！tf.boolean_mask的设计初衷就是提取符合条件的元素，所以必然会压缩维度，但你要的是保留原维度，将False位置置0，这其实是个掩码替换的需求，有两种简单高效的方法可以实现：

方法一：布尔掩码转数值型后与原张量相乘 #

布尔值在TensorFlow里可以直接转换为和原张量同类型的数值（True→1，False→0），和原张量逐元素相乘后，False对应的位置自然就变成0了，完美保留原维度：

import tensorflow as tf

# 你的输入
tensor = tf.constant([1,2,3,4,5])
bool_array = tf.constant([True, False, False, False, True])

# 转换掩码类型并相乘
result = tensor * tf.cast(bool_array, dtype=tensor.dtype)
print(result.numpy())  # 输出: [1 0 0 0 5]

如果你的布尔数组是numpy格式的，记得先转成TensorFlow张量再操作哦。

方法二：使用tf.where精准控制替换逻辑 #

tf.where可以根据条件选择两个张量中的元素，我们可以明确指定：当掩码为True时取原张量的值，为False时取0，逻辑更直观：

import tensorflow as tf

tensor = tf.constant([1,2,3,4,5])
bool_array = tf.constant([True, False, False, False, True])

result = tf.where(bool_array, tensor, tf.zeros_like(tensor))
print(result.numpy())  # 输出: [1 0 0 0 5]

这个方法可读性更强，能清晰看到两种情况的取值规则。

补充：为什么tf.boolean_mask不符合需求？ #

再解释下你疑惑的点：tf.boolean_mask(tensor, mask)的核心作用是筛选出mask中为True的元素，它会返回一个仅包含符合条件元素的张量（原例子里就是[1,5]），本质是维度压缩的提取操作，和你想要的“掩码替换并保留维度”不是同一个应用场景~

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