如何在TensorFlow中用掩码矩阵将一维张量赋值给二维张量？

在TensorFlow中，由于张量默认是不可变的，无法像PyTorch那样直接进行原地索引赋值操作，但可以通过以下两种常用方法实现需求：

方法一：使用tf.where（推荐，写法简洁） #

tf.where会根据掩码条件，在对应位置选择数组的值或零张量的值，直接生成目标张量：

import tensorflow as tf

# 定义所需张量
mask = tf.constant([True] * 11)  # 掩码（11个True）
array = tf.constant([1,2,3,4,5,6,7,8,9,10,11], dtype=tf.float32)  # 含11个元素的数组
zero_tensor = tf.zeros_like(mask, dtype=tf.float32)  # 同掩码形状的零张量

# 完成赋值操作
result_tensor = tf.where(mask, array, zero_tensor)

方法二：使用tf.scatter_nd #

通过获取掩码为True的位置索引，再将数组的值分散到零张量的对应位置：

import tensorflow as tf

# 定义所需张量
mask = tf.constant([True] * 11)
array = tf.constant([1,2,3,4,5,6,7,8,9,10,11], dtype=tf.float32)
zero_tensor = tf.zeros_like(mask, dtype=tf.float32)

# 获取掩码为True的位置索引
indices = tf.where(mask)
# 将数组值映射到零张量的对应位置
result_tensor = tf.scatter_nd(indices, array, shape=tf.shape(zero_tensor))

两种方法最终得到的result_tensor，都会在掩码为True的位置保留数组的对应元素，其余位置保持零张量的初始值，和PyTorch中zero_tensor[mask] = array的效果一致。

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