How to fix tf.function retracing in TensorFlow

Graph creation adds noticeable overhead when a frequently called TensorFlow function keeps missing its cached graph. Variable batch lengths are a common trigger because an unconstrained tf.function can specialize on every concrete input shape.

TensorFlow selects a cached ConcreteFunction from the tracing type of each argument. A tensor's shape and dtype contribute to that type, so batches shaped (1, 3) and (6, 3) can create separate traces even though both contain the same three features.

An explicit input_signature fixes the accepted rank, dtype, and feature width while None permits the batch dimension to vary. Inputs outside that contract fail at the call boundary instead of adding another graph specialization.

1. Create the shared TensorFlow input section at the start of tf_function_retracing.py.

   tf_function_retracing.py

   import os
    
   os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2"
    
   import tensorflow as tf
    
   tf.get_logger().setLevel("ERROR")
    
    
   def make_batch(rows):
       return tf.ones((rows, 3), dtype=tf.float32)

2. Add the unconstrained function and repeated batch calls after make_batch().

   @tf.function
   def score_batch(features):
       print(f"Tracing without a signature: {features.shape}")
       return tf.reduce_sum(features, axis=1)
    
    
   for rows in range(1, 7):
       score_batch(make_batch(rows))
    
   print(f"Trace count: {score_batch.experimental_get_tracing_count()}")

   Python print() runs while TensorFlow traces the function, so every printed shape represents a graph trace rather than an ordinary graph execution.

3. Run tf_function_retracing.py to reproduce shape-specific retracing.

   $ python3 tf_function_retracing.py
   Tracing without a signature: (1, 3)
   Tracing without a signature: (2, 3)
   Tracing without a signature: (3, 3)
   Tracing without a signature: (4, 3)
   Tracing without a signature: (5, 3)
   Tracing without a signature: (6, 3)
   Trace count: 6

4. Replace the unconstrained function and calls with a signature that permits variable batch lengths.

   @tf.function(
       input_signature=(
           tf.TensorSpec(shape=(None, 3), dtype=tf.float32, name="features"),
       )
   )
   def score_batch(features):
       print(f"Tracing with a signature: {features.shape}")
       return tf.reduce_sum(features, axis=1)
    
    
   for rows in range(1, 7):
       result = score_batch(make_batch(rows))
    
   print(f"Trace count: {score_batch.experimental_get_tracing_count()}")
   print(f"Final output shape: {result.shape}")
   print(f"Final output values: {result.numpy()}")

   Fixed feature dimensions and the expected dtype remain explicit; None is reserved for a dimension that the real function accepts at different lengths.

5. Rerun tf_function_retracing.py to confirm that all six batch lengths reuse one trace.

   $ python3 tf_function_retracing.py
   Tracing with a signature: (None, 3)
   Trace count: 1
   Final output shape: (6,)
   Final output values: [3. 3. 3. 3. 3. 3.]