TensorFlow中支持梯度计算的条件数组赋值实现方案咨询

已知条件：

• correct_idx：布尔数组
• outp2：与correct_idx形状相同的张量
• train_preds：张量，其元素数量等于correct_idx中True的个数

需求：实现与NumPy代码outp2[correct_idx] = train_preds等价的功能，且支持梯度计算。

之前尝试的代码因维度不匹配报错：

correct_idxs2  = tf.convert_to_tensor(correct_idx)
outp2          = tf.where(correct_idxs2, train_preds, tf.constant(float('nan'),dtype=tf.float32))

问题出在tf.where要求两个分支的张量形状必须和correct_idxs2完全一致，但train_preds的长度仅等于correct_idx中True的数量，和correct_idxs2的多维形状不匹配，因此触发维度异常。

可行方案：使用tf.tensor_scatter_nd_update，该操作支持梯度计算，能精准将train_preds的值填充到outp2中correct_idx为True的位置：

import tensorflow as tf

# 将布尔数组转为张量
correct_idxs_tensor = tf.convert_to_tensor(correct_idx)
# 获取所有为True的位置的索引
true_indices = tf.where(correct_idxs_tensor)
# 执行带梯度支持的赋值操作
updated_outp2 = tf.tensor_scatter_nd_update(
    tensor=outp2,
    indices=true_indices,
    updates=train_preds
)

这个方法的核心是先定位所有需要更新的位置索引，再通过tf.tensor_scatter_nd_update将train_preds的值逐个填充到对应位置，既保留了outp2中其他位置的原始值，又完全支持自动梯度计算。

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