The Story — The Tailor's Paper Pattern
The next customer has exactly the same size. The tailor does not measure again. He picks up the same pattern and starts cutting at once.
A customer with a different size arrives. Now the tailor must draw a new pattern and add it to his drawer.
@tf.function works the same way. Drawing the pattern is called tracing. The drawer of patterns is a cache of graphs. The "size" of the customer is the input signature.
Two Ways to Use tf.function
import tensorflow as tf
@tf.function # 1. as a decorator
def scale(x):
return x * 10
def shift(x):
return x + 1
fast_shift = tf.function(shift) # 2. wrap an existing function
print(scale(tf.constant([1.0, 2.0])))
print(fast_shift(tf.constant([1.0, 2.0])))Both give a polymorphic function: one Python name that can hold many graphs, one per input signature.
What Tracing Really Does
You can see tracing happen. Python's print runs only during tracing, and it sees a symbolic tensor. tf.print is a graph op, so it runs on every call with real values.
@tf.function
def double(x):
print(" [trace] Python sees:", x)
tf.print(" [run] value is", x)
return x * 2
print("call 1"); double(tf.constant(5))
print("call 2"); double(tf.constant(7))
print("call 3"); double(tf.constant(1.5))
print("traces:", double.experimental_get_tracing_count())Look Inside the Cache
print(double.pretty_printed_concrete_signatures())
double is a polymorphic function. It holds two concrete functions: one for int32 scalars and one for float32 scalars.
Get one directly with double.get_concrete_function(tf.TensorSpec([], tf.float32)).
Try It — The Trace Cache Simulator
Build an argument for f(x) and press Call. Watch the cache fill up. Turn on reduce_retracing and call with tensors of length 1, 2, 3 to see TensorFlow relax the shape to None.
| # | Concrete function signature | Times used |
|---|
AutoGraph — Python Control Flow Becomes Graph Ops
During tracing, a helper called AutoGraph rewrites your if, for and while statements. What it produces depends on one question: is the condition a tensor, or a Python value?
| Your code | Condition is a… | In the graph |
|---|---|---|
if x > 0: | Tensor | tf.cond — both branches are stored, one runs |
if flag: | Python bool | Branch chosen at trace time. The other is never in the graph. |
for i in tf.range(n): | Tensor | tf.while_loop — one small loop node |
for i in range(n): | Python int | Unrolled — the body is copied n times |
while x < 10: | Tensor | tf.while_loop |
@tf.function
def sign_flip(x):
if x > 0: # tensor condition -> tf.cond
return x
return -x
ops = [op.type for op in sign_flip.get_concrete_function(tf.constant(1)).graph.get_operations()]
print("tensor if :", ops)
@tf.function
def add_python_loop(x):
for _ in range(100): # Python loop -> unrolled
x = x + 1
return x
@tf.function
def add_tf_loop(x):
for _ in tf.range(100): # tensor loop -> tf.while_loop
x = x + 1
return x
for fn in (add_python_loop, add_tf_loop):
g = fn.get_concrete_function(tf.constant(0)).graph
print(f"{fn.__name__:16s}: {len(g.get_operations()):3d} ops, result {fn(tf.constant(0)).numpy()}")A Python range(10000) loop inside a tf.function copies the body ten thousand times. Tracing becomes very slow and memory grows. Loop over tf.range (or a tf.data.Dataset) when the loop is long.
Try It — The Loop Unroller
Choose the loop type and the number of steps. See how the traced graph grows. The op counts follow the real numbers from the code above (2n + 2 for an unrolled add loop).
Golden Rules
print runs at trace time only. tf.print runs on every call.tf.cond / tf.while_loop. Conditions on Python values are decided once, at trace time.tf.range for long loops inside a tf.function to avoid giant unrolled graphs.