TensorFlow中如何对满足条件的张量分段应用函数？

我有一个存储原点距离值的二维矩阵（将用于二维傅里叶变换）：

s = tf.linspace(0, 10, 100)
x_grid, y_grid = tf.meshgrid(s, s)

t = x_grid**2 + y_grid**2

我希望对该张量应用具有分段行为的TensorFlow函数：低于阈值的值使用fun_1，高于阈值的值使用fun_2。

在NumPy中，可通过以下方式轻松实现：

t[t <= threshold] = fun_1(t[t <= threshold])
t[t > threshold] = fun_2(t[t > threshold])

但在TensorFlow中执行相同操作会报错：TypeError: only integer scalar arrays can be converted to a scalar index。

我已查阅大量文档和资料仍未找到合适方案，请问有人解决过类似问题吗？

解决方案 #

TensorFlow的张量默认是不可变的，无法像NumPy那样直接原地索引赋值，以下几种方法可以实现需求：

方法1：使用tf.where（推荐，简洁高效）

tf.where会根据逐元素的条件，选择对应位置的输出结果：

# 确保fun_1和fun_2是兼容TensorFlow张量的函数
result = tf.where(t <= threshold, fun_1(t), fun_2(t))

该方法会对t的所有元素分别计算fun_1(t)和fun_2(t)，再根据条件筛选对应值，实现逐元素的分段逻辑。

方法2：结合掩码与tf.tensor_scatter_nd_update

如果fun_1或fun_2计算成本较高，不想对所有元素计算两个函数，可以用掩码筛选元素后再合并结果：

# 生成逐元素的布尔掩码
mask = t <= threshold

# 获取符合条件元素的索引
indices = tf.where(mask)

# 分别处理两部分元素
processed_part = fun_1(tf.gather_nd(t, indices))
base_part = fun_2(t)

# 将处理后的部分更新到基准张量中，得到最终结果
result = tf.tensor_scatter_nd_update(base_part, indices, processed_part)

注意：tf.tensor_scatter_nd_update返回的是新张量，不会修改原张量，符合TensorFlow的设计规范。

方法3：使用tf.cond（仅适用于全局条件判断）

如果你的需求是基于整个张量的全局判断（比如判断所有元素是否都小于阈值），可以用tf.cond，但不适用于逐元素的分段操作：

def apply_fun1():
    return fun_1(t)

def apply_fun2():
    return fun_2(t)

result = tf.cond(tf.reduce_all(t <= threshold), apply_fun1, apply_fun2)

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