Confused by the behavior of `tf.cond`

Question

I need a conditional control flow in my graph. If pred is True, the graph should call an op that updates a variable and then returns it, otherwise

Answer 1

TL;DR: If you want tf.cond() to perform a side effect (like an assignment) in one of the branches, you must create the op that performs the side effect inside the function that you pass to tf.cond().

The behavior of tf.cond() is a little unintuitive. Because execution in a TensorFlow graph flows forward through the graph, all operations that you refer to in either branch must execute before the conditional is evaluated. This means that both the true and the false branches receive a control dependency on the tf.assign() op, and so y always gets set to 2, even if pred isFalse`.

The solution is to create the tf.assign() op inside the function that defines the true branch. For example, you could structure your code as follows:

pred = tf.placeholder(tf.bool, shape=[])
x = tf.Variable([1])
def update_x_2():
  with tf.control_dependencies([tf.assign(x, [2])]):
    return tf.identity(x)
y = tf.cond(pred, update_x_2, lambda: tf.identity(x))
with tf.Session() as session:
  session.run(tf.initialize_all_variables())
  print(y.eval(feed_dict={pred: False}))  # ==> [1]
  print(y.eval(feed_dict={pred: True}))   # ==> [2]

Answer 2

pred = tf.constant(False)
x = tf.Variable([1])

def update_x_2():
    assign_x_2 = tf.assign(x, [2])
    with tf.control_dependencies([assign_x_2]):
        return tf.identity(x)
y = tf.cond(pred, update_x_2, lambda: tf.identity(x))
with tf.Session() as session:
  session.run(tf.initialize_all_variables())
  print(y.eval())

This will get the result of [1].

This answer is quite the same as the above answer. But what I wanna share is you can put every ops you would like to use in its branch function. Because, given your example code, tensor x is can be directly used by the update_x_2 function.