TensorFlow不支持布尔索引？如何实现类似NumPy的布尔索引功能

其实TensorFlow（尤其是2.x版本）是支持布尔索引的，只是因为TensorFlow张量的特性（默认不可变，除非用tf.Variable），用法和NumPy有一点点差异。下面我会对应你给出的NumPy示例，给出TensorFlow的实现方式：

方法一：使用tf.Variable直接赋值（对应NumPy的原地修改） #

普通的tf.Tensor是不可变的，所以我们需要用tf.Variable创建可修改的张量，之后就能像NumPy那样用布尔索引赋值了：

import tensorflow as tf

# 创建可修改的变量张量
A = tf.Variable([3, 4, 5, -1, 6, -1, 7, 8], dtype=tf.int32)
mask = (A == -1)
print("原始张量：", A.numpy())

# 布尔索引赋值
A[mask].assign([11, 12])
print("修改后的张量：", A.numpy())

运行这段代码的输出和你的NumPy示例完全一致：

原始张量： [3 4 5 -1 6 -1 7 8]
修改后的张量： [ 3 4 5 11 6 12 7 8]

方法二：使用tf.tensor_scatter_nd_update生成新张量（不原地修改） #

如果你不想使用可变的tf.Variable，可以通过生成新张量的方式实现逻辑：

import tensorflow as tf

A = tf.constant([3, 4, 5, -1, 6, -1, 7, 8], dtype=tf.int32)
mask = (A == -1)
# 获取满足条件的元素索引
indices = tf.where(mask)
# 准备替换的值
replace_values = tf.constant([11, 12], dtype=tf.int32)
# 生成新的修改后张量
updated_A = tf.tensor_scatter_nd_update(A, indices, replace_values)

print("原始张量：", A.numpy())
print("修改后的张量：", updated_A.numpy())

这个方法适合处理不可变的常量张量，同样能得到你想要的结果。

额外小技巧 #

布尔索引不仅可以用来赋值，还能提取符合条件的元素：用tf.boolean_mask(A, mask)可以直接取出所有满足条件的元素，和NumPy里A[mask]提取元素的用法完全一致。

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