Introduction

A boolean tensor looks like something you should be able to drop directly into a normal Python if statement, but that is usually the wrong mental model. In TensorFlow, tensors are symbolic or deferred values in many execution paths, so Python control flow and tensor control flow are not always interchangeable. The right tool depends on whether you have one scalar condition or an element-wise mask.

The short answer is this: use tf.cond when one boolean tensor decides between two branches, and use tf.where when you want per-element selection across tensors.

Use tf.cond for a Single Tensor Condition

If your condition is a scalar boolean tensor, tf.cond is the direct equivalent of an if statement in graph-style execution.

1import tensorflow as tf
2
3flag = tf.constant(True)
4value = tf.constant(10)
5
6result = tf.cond(
7    flag,
8    lambda: value + 1,
9    lambda: value - 1,
10)
11
12print(result.numpy())

This says: if flag is true, run the first branch; otherwise run the second branch. Each branch is wrapped in a zero-argument function so TensorFlow can build or execute the correct branch at the right time.

That is the main pattern to remember when you want one tensor to decide one code path.

Why a Plain Python if Usually Fails

A common first attempt looks like this:

1import tensorflow as tf
2
3flag = tf.constant(True)
4
5if flag:
6    print("yes")

That is fragile because Python wants an immediate boolean value, while a tensor is a TensorFlow object. In eager mode you may still run into type errors, and inside @tf.function the difference becomes even more important because TensorFlow is tracing code.

If you are writing TensorFlow logic, treat tensor conditions as tensor control flow, not as ordinary Python booleans.

Use tf.where for Element-Wise Selection

Sometimes you do not want a single branch for the whole program. You want a boolean tensor to choose values position by position. That is what tf.where does.

1import tensorflow as tf
2
3mask = tf.constant([True, False, True, False])
4a = tf.constant([10, 20, 30, 40])
5b = tf.constant([1, 2, 3, 4])
6
7result = tf.where(mask, a, b)
8print(result.numpy())

The output uses elements from a where the mask is true and elements from b where the mask is false.

This is not the same as tf.cond. tf.cond chooses one branch for the whole expression. tf.where chooses values element by element.

Python if Can Still Appear Inside @tf.function

TensorFlow can convert some Python control flow automatically when the code is wrapped in @tf.function. That can make this pattern work:

1import tensorflow as tf
2
3@tf.function
4def adjust(x):
5    if tf.reduce_mean(x) > 0:
6        return x * 2
7    return x - 2
8
9print(adjust(tf.constant([1.0, 2.0])).numpy())

This works because TensorFlow transforms supported Python control flow into graph-compatible operations during tracing. Even so, it is still valuable to understand tf.cond, because it makes the intent explicit and behaves well when you need a direct tensor-branching primitive.

Build Clear Branches

When using tf.cond, both branches should return compatible tensor shapes and types. A small helper function often keeps the code readable.

1import tensorflow as tf
2
3def choose_loss(use_logits, y_true, y_pred):
4    return tf.cond(
5        use_logits,
6        lambda: tf.nn.sigmoid_cross_entropy_with_logits(labels=y_true, logits=y_pred),
7        lambda: tf.keras.losses.binary_crossentropy(y_true, y_pred),
8    )

The key idea is that both branches represent valid tensor computations. If one branch returns a scalar and the other returns a rank-two tensor, the result becomes hard to reason about and may fail.

Common Pitfalls

The most common mistake is trying to use a tensor as though it were an ordinary Python boolean. That usually leads to confusing control-flow errors.

Another frequent mistake is using tf.cond when the real goal is element-wise masking. In that case, tf.where is the better tool.

A third issue is returning incompatible shapes or types from the two branches of tf.cond. Even if the code looks symmetric, the outputs still need to line up.

Summary

• Use tf.cond when one boolean tensor chooses between two whole branches.
• Use tf.where when a boolean tensor should select values element by element.
• Do not rely on a plain Python if with tensors unless you understand how TensorFlow tracing applies.
• Keep tf.cond branch outputs compatible in shape and dtype.
• Think in tensor control flow, not ordinary Python boolean semantics.