How to fix tf.function retracing in TensorFlow

tf.function builds graph-backed TensorFlow callables, but graph creation has its own cost. Repeated retracing usually appears when a hot training or inference path keeps calling the same decorated function with new shapes, dtypes, or Python values, so the runtime spends time creating graphs instead of executing the cached one.

TensorFlow caches ConcreteFunction graphs behind each tf.function and selects one by the call signature. Tensor shape, dtype, and Python-valued arguments all participate in that signature, while an input_signature can tell TensorFlow which tensor dimensions are allowed to vary without creating another graph.

Start by proving which call pattern retraces, then make the function contract stable. Use input_signature when the tensor rank, dtype, and non-batch dimensions are known, convert changing numeric options to tensors when they do not change graph structure, and reserve reduce_retracing=True for functions that need automatic shape generalization.

1. Save a retracing diagnostic as

   tf_function_retracing_demo.py

   .

   tf_function_retracing_demo.py

   import os
    
   os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2"
    
   import tensorflow as tf
    
   tf.get_logger().setLevel("WARNING")
    
    
   def sample_batches():
       return [tf.ones((rows, 3), dtype=tf.float32) for rows in range(1, 8)]
    
    
   @tf.function
   def score_unfixed(features):
       print(f"Tracing unfixed graph for shape {features.shape}")
       return tf.reduce_sum(features, axis=1)
    
    
   @tf.function(
       input_signature=[
           tf.TensorSpec(shape=(None, 3), dtype=tf.float32, name="features")
       ]
   )
   def score_fixed(features):
       print(f"Tracing fixed graph for shape {features.shape}")
       return tf.reduce_sum(features, axis=1)
    
    
   print(f"TensorFlow {tf.__version__}")
   print("Unfixed calls:")
   for batch in sample_batches():
       score_unfixed(batch)
   print(f"Unfixed tracing count: {score_unfixed.experimental_get_tracing_count()}")
    
   print("")
   print("Fixed calls:")
   last_output = None
   for batch in sample_batches():
       last_output = score_fixed(batch)
   print(f"Fixed tracing count: {score_fixed.experimental_get_tracing_count()}")
   print(f"Fixed last output shape: {tuple(last_output.shape.as_list())}")

   The Python print() calls are trace markers for the diagnostic. They run when TensorFlow traces a graph, not every time the cached graph executes.

2. Run the diagnostic and confirm that changing only the batch length triggers repeated traces without a signature.

   $ python3 tf_function_retracing_demo.py
   TensorFlow 2.21.0
   Unfixed calls:
   Tracing unfixed graph for shape (1, 3)
   Tracing unfixed graph for shape (2, 3)
   Tracing unfixed graph for shape (3, 3)
   Tracing unfixed graph for shape (4, 3)
   Tracing unfixed graph for shape (5, 3)
   WARNING:tensorflow:5 out of the last 5 calls to <function score_unfixed> triggered tf.function retracing.
   ##### snipped #####
   Tracing unfixed graph for shape (6, 3)
   WARNING:tensorflow:6 out of the last 6 calls to <function score_unfixed> triggered tf.function retracing.
   ##### snipped #####
   Tracing unfixed graph for shape (7, 3)
   Unfixed tracing count: 7
   
   Fixed calls:
   Tracing fixed graph for shape (None, 3)
   Fixed tracing count: 1
   Fixed last output shape: (7,)

   The fixed function accepts different batch lengths through shape=(None, 3) and still keeps one trace.

3. Add an input_signature to the real function whose tensor shape varies by batch size.

   @tf.function(
       input_signature=[
           tf.TensorSpec(shape=(None, 3), dtype=tf.float32, name="features")
       ]
   )
   def score_batch(features):
       return tf.reduce_sum(features, axis=1)

   Use None only for dimensions that may vary between calls. Keep the dtype and fixed feature dimensions explicit so incompatible inputs fail early instead of creating extra traces.

   Related: How to compile a function with tf.function in TensorFlow

4. Pass changing numeric options as tensors when they do not change the graph structure.

   @tf.function
   def scale_scores(scores, temperature):
       return scores / temperature
    
   scores = tf.constant([0.35, 0.80], dtype=tf.float32)
    
   scale_scores(scores, tf.constant(0.7, dtype=tf.float32))
   scale_scores(scores, tf.constant(0.9, dtype=tf.float32))

   Keep Python booleans, strings, or integers only when the value intentionally selects different graph-building logic.

5. Keep the decorated callable outside the repeated training, inference, or request loop.

   @tf.function(
       input_signature=[
           tf.TensorSpec(shape=(None, 3), dtype=tf.float32, name="features")
       ]
   )
   def train_step(features):
       return tf.reduce_sum(features, axis=1)
    
   for features in dataset:
       train_step(features)

   Creating a new tf.function inside a loop defeats the cache because each decorated callable owns its own trace cache.

   Related: How to run a custom training loop in TensorFlow

6. Use reduce_retracing=True only when a single explicit signature is too narrow for the supported inputs.

   @tf.function(reduce_retracing=True)
   def normalize(values):
       return values - tf.reduce_mean(values)

   reduce_retracing=True lets TensorFlow generalize observed tensor shapes automatically. An explicit input_signature is still the tighter contract when the allowed shape and dtype are known.

7. Rerun the affected workload and confirm that the retracing warning no longer repeats.

   $ python3 tf_function_retracing_demo.py
   ##### snipped #####
   Fixed calls:
   Tracing fixed graph for shape (None, 3)
   Fixed tracing count: 1
   Fixed last output shape: (7,)

   If the warning remains, check for changing Python objects, new dtypes, or a tf.function decorator created inside the loop before widening the signature further.