Introduction

TensorFlow graph execution has different control flow rules than regular Python scripts. Conditional branches and loops must be written with TensorFlow ops when you want code to run correctly under tf.function. If you mix Python control flow with dynamic tensor shapes, you can get tracing errors or incorrect behavior.

Conditional Graph Logic with tf.cond

In eager mode, a normal Python if works for immediate values. Inside graph tracing, conditions that depend on tensors should use tf.cond so both branches are represented in the graph and selected at runtime.

1import tensorflow as tf
2
3@tf.function
4def threshold_scale(x, threshold):
5    mean_val = tf.reduce_mean(x)
6
7    def high_branch():
8        return x * 0.5
9
10    def low_branch():
11        return x * 2.0
12
13    return tf.cond(mean_val > threshold, high_branch, low_branch)
14
15x1 = tf.constant([1.0, 2.0, 3.0])
16x2 = tf.constant([10.0, 12.0, 14.0])
17
18print(threshold_scale(x1, tf.constant(5.0)))
19print(threshold_scale(x2, tf.constant(5.0)))

This pattern is useful when model paths differ by tensor statistics, sequence lengths, or confidence scores.

Looping with Tensor Sizes

Using len(tensor) in graph mode is unsafe for dynamic shapes. Prefer tf.shape(tensor) and TensorFlow loop constructs. For simple reductions, vectorized ops are best. When you truly need iteration, use tf.range or tf.while_loop.

1import tensorflow as tf
2
3@tf.function
4def row_sums(matrix):
5    n_rows = tf.shape(matrix)[0]
6    out = tf.TensorArray(dtype=matrix.dtype, size=n_rows)
7
8    for i in tf.range(n_rows):
9        row_total = tf.reduce_sum(matrix[i])
10        out = out.write(i, row_total)
11
12    return out.stack()
13
14m = tf.constant([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]])
15print(row_sums(m))

For variable length loops with changing state, tf.while_loop gives explicit control and often cleaner graph behavior.

1import tensorflow as tf
2
3@tf.function
4def cumulative_until_limit(values, limit):
5    n = tf.shape(values)[0]
6
7    def cond(i, total):
8        return tf.logical_and(i < n, total < limit)
9
10    def body(i, total):
11        return i + 1, total + values[i]
12
13    i0 = tf.constant(0)
14    total0 = tf.constant(0.0)
15    _, total = tf.while_loop(cond, body, [i0, total0])
16    return total
17
18vals = tf.constant([0.7, 1.2, 2.3, 4.0], dtype=tf.float32)
19print(cumulative_until_limit(vals, tf.constant(3.0)))

Performance and Readability Tradeoffs

When possible, replace loops with vectorized operations such as tf.reduce_sum, tf.map_fn, or broadcasting. Vectorized code is usually faster and easier to optimize on accelerators.

Still, explicit graph loops are valid when operations are stateful, early stopping is required, or each step uses different branch logic. The key is consistency: keep graph compatible ops together and avoid Python side effects inside traced functions.

If a function retraces too often, check whether input shapes and dtypes vary across calls. Constraining signatures can reduce retracing cost in production.

Testing control flow functions with both small and large tensors is important. Some bugs only appear with edge sizes such as empty tensors or one element sequences. Add unit tests for these cases so graph behavior stays stable during model refactoring.

Common Pitfalls

A common mistake is writing if tensor_value: inside tf.function. Tensor truth checks are ambiguous in graph mode and typically fail. Use tf.cond or tensor comparisons that feed TensorFlow control flow ops.

Another issue is using Python range with tensor length from tf.shape. Python range needs a concrete integer, but dynamic tensor shapes are symbolic during tracing. Use tf.range for graph friendly loops.

Developers also create Python lists inside traced loops and append tensors on each iteration. This can break graph conversion. Use tf.TensorArray for accumulation.

Finally, mixing int32 and float32 in loop state often causes type errors. Initialize loop variables with the exact dtypes you need and keep them consistent across loop body returns.

Summary

• Use tf.cond for tensor dependent branching inside graph execution.
• Use tf.shape and tf.range instead of Python length and loops.
• Prefer vectorized ops, but use graph loops when control flow requires them.
• Accumulate loop outputs with tf.TensorArray in traced functions.
• Keep dtypes consistent to avoid hard to debug graph errors.