Example of tf.cond usage to allow gradient propagation through conditional branches

tf_cond_example_gradients.py

import tensorflow as tf

# X, Y and Z are all watched variables, although they will produce a gradient (i.e. not None) only if the branch executed has uses that variable
with tf.GradientTape() as tape:
    x = tf.Variable(5.0, dtype=tf.float32, name='X')
    y = tf.Variable(6.0, dtype=tf.float32, name='Y')
    z = tf.Variable(8.0, dtype=tf.float32, name='Z')

    cond = tf.cond(pred = x < y,
               true_fn = lambda: tf.add(x, z),
               false_fn = lambda: tf.square(y))
    print(tape.watched_variables())
    op = tape.gradient(cond, tape.watched_variables()) 
    print(op)

# in this case, only X and Y are watched and X will have gradient=None
with tf.GradientTape() as tape:
    x = tf.Variable(5.0, dtype=tf.float32, name='X')
    y = tf.Variable(6.0, dtype=tf.float32, name='Y')
    z = tf.Variable(8.0, dtype=tf.float32, name='Z')

    cond = tf.cond(pred = x > y,
               true_fn = lambda: tf.add(y, y),
               false_fn = lambda: tf.square(y))
    print(tape.watched_variables())
    op = tape.gradient(cond, tape.watched_variables()) 
    print(op)