Introduction

Yes. In TensorFlow graph mode, tf.gradients can propagate through tf.cond as long as the executed branch is differentiable with respect to the input you care about. The important detail is that the gradient follows the branch that actually runs, not both branches equally.

tf.cond Chooses One Branch at Runtime

tf.cond(pred, true_fn, false_fn) builds a graph with two possible subgraphs and selects one of them at execution time based on pred. From the gradient engine's perspective, that means TensorFlow differentiates through the operations that contributed to the chosen output.

1import tensorflow as tf
2
3tf.compat.v1.disable_eager_execution()
4
5x = tf.compat.v1.placeholder(tf.float32, shape=())
6
7y = tf.cond(
8    x > 0,
9    lambda: x * x,
10    lambda: 3.0 * x
11)
12
13grad = tf.gradients(y, x)[0]
14
15with tf.compat.v1.Session() as sess:
16    print(sess.run([y, grad], feed_dict={x: 2.0}))
17    print(sess.run([y, grad], feed_dict={x: -2.0}))

For x = 2.0, the active branch is x * x, so the gradient is 2x, which is 4.0. For x = -2.0, the active branch is 3x, so the gradient is 3.0.

The Inactive Branch Does Not Drive the Gradient

Developers sometimes imagine that TensorFlow differentiates both branches and blends the results. That is not what happens here. The branch not taken does not contribute to the current output value, so it does not contribute to the current gradient path either.

This makes tf.cond behave more like real control flow than elementwise selection. If you need branch-like behavior at the element level inside tensors, tf.where may be the more relevant operator, but it has different semantics and gradient behavior.

Both Branches Still Need Compatible Outputs

Even though only one branch runs at execution time, both branches must still return compatible tensors when the graph is built. In practice that means matching dtype and compatible shape.

For example, this is conceptually valid:

1y = tf.cond(
2    x > 0,
3    lambda: x * x,
4    lambda: x + 5.0
5)

Both branches return scalar float32 tensors. If one branch returned a scalar and the other returned a vector, the graph would be ill-formed for most normal uses.

Differentiability Still Matters

tf.gradients will pass through tf.cond only if the active branch contains differentiable operations all the way back to the target input. If the chosen branch uses an op without a defined gradient, the gradient calculation can fail or return None depending on the graph and requested tensors.

That means the practical rule is:

• 'tf.cond itself does not block gradients.'
• Non-differentiable ops inside a branch can still block them.

This is the same principle you would apply outside tf.cond. The conditional is not the real problem; the ops inside it might be.

Piecewise Functions Are a Common Use Case

Conditional branches are often used for piecewise mathematical definitions, clipping logic, or loss functions that change formula in different regions. In those cases, the gradient is piecewise too.

That is usually fine, but be careful near the branch boundary. At the exact switching point, the function may be non-smooth or ambiguous from an optimization perspective even if TensorFlow can still construct a graph.

TensorFlow 2 Perspective

In TensorFlow 2, you are more likely to use tf.GradientTape than tf.gradients, but the high-level idea is the same: differentiation follows the executed branch if the operations are differentiable. The title question is rooted in TensorFlow 1 style APIs, but the underlying principle carries over.

Common Pitfalls

• Assuming tf.cond blocks gradients by itself.
• Forgetting that only the executed branch contributes to the gradient.
• Returning incompatible shapes or dtypes from the two branches.
• Using non-differentiable operations inside the active branch and blaming tf.cond.
• Confusing tf.cond control flow with elementwise selection semantics from tf.where.

Summary

• 'tf.gradients can propagate through tf.cond.'
• The gradient follows the branch that actually executes.
• Both branches must still return compatible tensor structures.
• Differentiability depends on the ops inside the chosen branch.
• The same idea applies in TensorFlow 2 with GradientTape.