如何基于条件修改Tensor值？解决TensorFlow类NumPy索引赋值报错问题

报错原因 #

TensorFlow 不支持 NumPy 风格的布尔索引原地赋值语法，你使用的class_max[class_max == 1] *= class_weights[0]写法在静态计算图编译阶段无法稳定推导Mask的维度，当计算图无法确认Mask不是标量时就会触发该报错。同时TensorFlow的普通Tensor默认是不可变对象，本身也不支持直接的原地修改操作。

可用解决方案 #

方案1：使用tf.where实现（推荐，通用性最强） #

这是最简洁也最兼容的写法，同时支持动态图、静态图以及@tf.function装饰的函数场景：

# 满足条件的位置执行乘法，不满足条件的位置保留原始值
class_max = tf.where(
    condition = class_max == 1,
    x = class_max * class_weights[0],
    y = class_max
)

如果class_max是tf.Variable类型，可以配合assign方法使用：

class_max.assign(tf.where(class_max == 1, class_max * class_weights[0], class_max))

方案2：使用tf.tensor_scatter_nd_update实现 #

适合多条件复杂更新的场景，逻辑更灵活：

# 生成布尔Mask
mask = class_max == 1
# 提取满足条件的元素索引
indices = tf.where(mask)
# 计算更新后的值
updated_values = tf.gather_nd(class_max, indices) * class_weights[0]
# 赋值得到新的Tensor
class_max = tf.tensor_scatter_nd_update(class_max, indices, updated_values)

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