如何在TensorFlow中对张量进行带填充的元素掩码操作？

给定以下TensorFlow张量：

>>> input = tf.random.normal([2,3,5])
>>> input
<tf.Tensor: shape=(2, 3, 5), dtype=float32, numpy=
array([[[ 1.1260294 , -0.05932725,  0.85893923, -1.5332409 ,
          0.6681451 ],
        [ 0.8833729 ,  0.8421117 , -0.60990584,  0.08593109,
          0.5969471 ],
        [ 0.20015325, -0.9459327 , -1.0818844 , -1.7254639 ,
         -0.51545954]],

       [[-0.36073774, -0.24315724,  1.5217028 ,  1.5075827 ,
          0.05745999],
        [-0.2570101 ,  1.5501927 , -0.17113225,  0.16063859,
         -0.95638955],
        [ 0.48955616,  0.11943919, -0.3523262 ,  0.10750653,
          1.1027677 ]]], dtype=float32)>

>>> mask = tf.constant([[0,1,0],[1,0,1]])
>>> mask
<tf.Tensor: shape=(2, 3), dtype=int32, numpy=
array([[0, 1, 0],
       [1, 0, 1]], dtype=int32)>

需求是：根据mask中值为1的位置保留input对应元素，将mask为0的位置替换为0，同时保持输出张量形状为(2,3,5)，最终输出如下：

>>> masked_input
<tf.Tensor: shape=(2, 3, 5), dtype=float32, numpy=
array([[[ 0.8833729 ,  0.8421117 , -0.60990584,  0.08593109,
          0.5969471 ],
        [ 0 , 0 , 0 , 0 ,
         0],
        [ 0 , 0 , 0 , 0 ,
         0]],

       [[-0.36073774, -0.24315724,  1.5217028 ,  1.5075827 ,
          0.05745999],
        [ 0.48955616,  0.11943919, -0.3523262 ,  0.10750653,
          1.1027677 ],
        [ 0 , 0 , 0 , 0 ,
          0]]], dtype=float32)>

已尝试的方法存在局限：

• tf.gather：未找到合适的实现方式
• tf.boolean_mask：会丢弃维度，无法保持原形状
• tf.ragged.boolean_mask：返回不规则张量，不符合要求

可以通过以下步骤实现需求：

1. 扩展掩码维度：将mask从(2,3)扩展为(2,3,1)，使其与input的最后一维匹配，便于索引和广播运算。
2. 提取有效元素：定位mask为1的位置，收集对应input中的元素。
3. 零填充至原形状：将有效元素按顺序填充到初始全零张量中，剩余位置保持为0。

具体代码实现：

import tensorflow as tf

# 定义输入和掩码
input = tf.random.normal([2,3,5])
mask = tf.constant([[0,1,0],[1,0,1]])

# 扩展掩码维度，匹配输入张量的最后一维
mask_expanded = tf.expand_dims(mask, axis=-1)

# 获取所有mask为1的位置索引
valid_indices = tf.where(tf.cast(mask_expanded, tf.bool))

# 收集输入张量中的有效元素
valid_elements = tf.gather_nd(input, valid_indices)

# 统计每个样本需要保留的元素数量
per_batch_counts = tf.reduce_sum(mask, axis=1)

# 创建初始全零张量
masked_input = tf.zeros_like(input)

# 构造填充位置的索引：每个样本的有效元素对应张量的前N个位置（N为该样本的有效元素数）
batch_indices = tf.repeat(tf.range(tf.shape(input)[0]), per_batch_counts)
seq_indices = tf.concat([tf.range(count) for count in per_batch_counts], axis=0)
fill_indices = tf.stack([batch_indices, seq_indices], axis=1)

# 将有效元素填充到全零张量中
masked_input = tf.tensor_scatter_nd_update(masked_input, fill_indices, valid_elements)

print(masked_input)

代码说明：

• tf.expand_dims：扩展掩码维度，确保能与输入张量的维度对齐，方便后续索引操作。
• tf.where：精准定位所有掩码为1的元素位置，返回三维索引。
• tf.gather_nd：根据索引批量提取输入中的有效元素。
• tf.reduce_sum：统计每个样本内需要保留的元素数量，为后续填充位置提供依据。
• tf.tensor_scatter_nd_update：将有效元素填充到全零张量的指定位置，实现保留有效元素+零填充的效果。

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