TensorFlow：如何根据布尔掩码填充一维张量生成指定形状张量？

核心实现逻辑 #

依托TensorFlow的张量索引更新能力，将一维张量的元素按掩码真值的排列顺序，填入同形状的零张量对应位置，非真值位置保留0值即可得到目标输出。

可直接运行的实现代码 #

import tensorflow as tf

# 定义输入张量
a = tf.constant([0.1, 0.2, 0.3, 0.4])
b = tf.constant([[1, 1, 0], [0, 1, 1]])

# 1. 将数值型掩码转为布尔类型，标记待填充的位置
b_bool = tf.cast(b, dtype=tf.bool)

# 2. 初始化和掩码形状一致、和a数据类型匹配的全零张量作为输出基底
output = tf.zeros_like(b, dtype=a.dtype)

# 3. 提取掩码中真值的坐标，将a的数值按顺序填充到对应位置
true_coords = tf.where(b_bool)
output = tf.tensor_scatter_nd_update(output, indices=true_coords, updates=a)

运行结果 #

执行上述代码得到的输出张量完全匹配需求：

[[0.1, 0.2, 0.0],
 [0.0, 0.3, 0.4]]

注意事项 #

• 初始化全零张量时必须指定和a一致的数值类型，否则会因为整数、浮点类型不兼容直接触发报错
• 可提前添加校验逻辑避免输入不匹配的问题：通过tf.reduce_sum(tf.cast(b_bool, tf.int32)) == tf.shape(a)[0]判断掩码真值总数和a的长度是否一致，不一致时提前抛错，避免出现难排查的索引越界问题
• tf.where返回的真值坐标默认是按行优先顺序排列的，只要掩码真值总数和a的长度完全相等，就能保证填充顺序和示例要求的输出完全匹配，无需额外调整排序规则

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