如何用选中索引将TensorFlow中的Rank-2张量指定行置零？

当然可行啦！在TensorFlow里轻松就能实现这个需求，我给你两种实用的实现方式，你可以按需选择：

方法一：利用张量散射更新构建掩码 #

这种方法通过创建一个掩码张量，把需要置零的行对应的掩码设为0，再和原张量相乘即可：

import tensorflow as tf

# 定义示例张量A和选中索引B
A = tf.constant([[1,1,1,1], [2,2,2,2], [3,3,3,3], [4,4,4,4], [5,5,5,5]])
B = tf.constant([1, 3, 4])

# 创建和A行数一致的列向量掩码，初始全为1
mask = tf.ones([tf.shape(A)[0], 1], dtype=A.dtype)
# 将B中索引对应的掩码位置更新为0
mask = tf.tensor_scatter_nd_update(
    mask, 
    tf.expand_dims(B, axis=1),  # 把B转为[[1], [3], [4]]的形状，符合散射更新的索引格式
    tf.zeros([len(B), 1], dtype=A.dtype)
)

# 元素相乘，实现指定行置零
result = A * mask

# 查看结果
print(result.numpy())

运行后输出就是你想要的效果：

[[1 1 1 1]
 [0 0 0 0]
 [3 3 3 3]
 [0 0 0 0]
 [0 0 0 0]]

方法二：使用布尔掩码实现 #

这种方法通过判断每个行索引是否在B中，生成布尔掩码后转成数值掩码，再和原张量相乘：

import tensorflow as tf

A = tf.constant([[1,1,1,1], [2,2,2,2], [3,3,3,3], [4,4,4,4], [5,5,5,5]])
B = tf.constant([1, 3, 4])

# 生成布尔掩码：判断每个行索引是否不在B中（True表示保留，False表示置零）
bool_mask = tf.math.logical_not(tf.math.in1d(tf.range(tf.shape(A)[0]), B))
# 将布尔掩码转为数值类型，并扩展为列向量以便和A广播相乘
num_mask = tf.cast(tf.expand_dims(bool_mask, axis=1), dtype=A.dtype)

# 计算结果
result = A * num_mask

print(result.numpy())

这个方法和第一种效果完全一致，逻辑更直观，适合喜欢用布尔判断的场景。

两种方法的核心思路都是通过掩码实现对指定行的过滤，本质上都是利用TensorFlow的广播机制，让掩码和原张量逐元素相乘，从而把目标行置零。

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