Topolearntf.function is where TF's speed and exportability come from, and misunderstanding tracing causes its most confusing bugs.
Decorating a function with `@tf.function` turns it into a callable that builds a graph the first time it sees a given input signature (tracing) and then runs that graph on later calls. Graphs allow optimisation, parallel op execution, XLA compilation and export to SavedModel.
Tracing runs your Python code once, with symbolic tensors. Consequences:
- Python side effects (
print, appending to a list, incrementing a Python counter) happen only while tracing, not on every call. Usetf.printfor output that should happen each time. - Retracing happens when the signature changes: a new tensor shape or dtype, or a different Python value for a non-tensor argument. Passing Python numbers as arguments (e.g. a step count) triggers a retrace for every new value; pass tensors instead.
- Control flow on tensors is converted to graph ops by AutoGraph, while control flow on Python values is unrolled at trace time.
Limit retracing by giving an input_signature with tf.TensorSpec([None], tf.float32) (None = any size) or reduce_retracing=True. Check the count with f.experimental_get_tracing_count().
While debugging, tf.config.run_functions_eagerly(True) runs decorated functions as plain Python.
Example
import tensorflow as tf
@tf.function
def f(x):
print("tracing with shape", x.shape) # Python side effect: tracing only
tf.print("running") # graph op: every call
return x * 2
f(tf.constant([1.0, 2.0])) # tracing with shape (2,) / running
f(tf.constant([3.0, 4.0])) # running (same signature, graph reused)
f(tf.constant([1.0, 2.0, 3.0])) # traces again for the new shape, then running
print(f.experimental_get_tracing_count()) # 2
@tf.function(input_signature=[tf.TensorSpec([None], tf.float32)])
def total(x):
return tf.reduce_sum(x)
print(total(tf.constant([1.0, 2.0])).numpy()) # 3.0
print(total(tf.constant([1.0, 2.0, 3.0])).numpy()) # 6.0, no retraceRemember
- First call per signature traces a graph; later calls reuse it
- Python side effects run only during tracing; use
tf.print - Pass tensors, not Python numbers, to avoid retracing
Watch out for
- Calling a tf.function in a loop with a Python int argument, which retraces every iteration.