python - 在 TensorFlow 中编写分段函数/在 TensorFlow 中编写 if then

Question

如何编写分段 TensorFlow 函数，即其中包含 if 语句的函数？

当前代码

import tensorflow as tf

my_fn = lambda x : x ** 2 if x > 0 else x + 5
with tf.Session() as sess:
  x = tf.Variable(tf.random_normal([100, 1]))
  output = tf.map_fn(my_fn, x)

错误:

TypeError:不允许使用 tf.Tensor 作为 Python bool。使用 if t is not None: 而不是 if t: 来测试是否定义了张量，并使用逻辑 TensorFlow 操作来测试张量的值。

Answer 1

tf.select 也不再像该线程所指示的那样工作 https://github.com/tensorflow/tensorflow/issues/8647

对我有用的是 tf.where

condition = tf.greater(x, 0)
res = tf.where(condition, tf.square(x), x + 5)