AutoGraph and graph control flow

Knowing what AutoGraph converts explains slow traces, shape errors in loops and many tf.function surprises.

When tf.function traces your code, AutoGraph first rewrites its Python source. Control-flow statements are converted to functions that decide at trace time what to do:

If the condition is a tensor, if becomes tf.cond and while becomes tf.while_loop: the graph contains the control flow, evaluated at run time.
If the condition is a Python value, the statement runs normally during tracing: an if picks one branch, and for i in range(3) is unrolled into three copies of the body.
for x in tf.range(n) or over a tf.data.Dataset becomes a graph loop.

Graph loops impose rules. Variables modified in a loop must keep the same dtype and a compatible shape across iterations (use tf.autograph.experimental.set_loop_options(shape_invariants=...) when a shape legitimately changes). Both branches of a tensor if must define the same variables with compatible types.

Appending to a Python list in a graph loop does not work; use `tf.TensorArray` and stack() it at the end.

Large Python-unrolled loops produce huge graphs and slow tracing, a common hidden cost. Use tf.autograph.to_code(fn) to see the generated code, and remember that tf.Variables may only be created on the first trace.

Example

import tensorflow as tf
@tf.function
def collatz_steps(n):
    steps = tf.constant(0)
    while n != 1:              # tensor condition -> tf.while_loop
        if n % 2 == 0:         # tensor condition -> tf.cond
            n = n // 2
        else:
            n = 3 * n + 1
        steps += 1
    return steps

print(collatz_steps(tf.constant(6)).numpy())    # 8
@tf.function
def loops(x):
    for _ in range(3):         # Python range: unrolled at trace time
        x = x * 2
    acc = tf.TensorArray(tf.float32, size=0, dynamic_size=True)
    for i in tf.range(3):      # tensor range: a real graph loop
        acc = acc.write(i, tf.cast(i, tf.float32) * x)
    return acc.stack()

print(loops(tf.constant(1.0)).numpy())          # [ 0.  8. 16.]
print(tf.autograph.to_code(collatz_steps.python_function)[:60])

Remember

Tensor conditions become tf.cond / tf.while_loop; Python conditions run at trace time
range() loops unroll; tf.range() loops stay in the graph
Accumulate in loops with tf.TensorArray

Watch out for

Appending tensors to a Python list inside a tf.range loop, which fails or captures only the trace-time value.

Check yourself

Inside a tf.function, what happens to `for i in range(1000): x = x + 1`?

Official docs: www.tensorflow.org/guide/function

Example

Remember

Watch out for

Check yourself

Builds on

Unlocks