Why does a tf.function retrace, and how do you stop it retracing on every call?

## The cache key A `tf.function` is not one graph. It is a polymorphic callable holding a dictionary of graphs, and the key is the **input signature** of the call. Understanding what goes into that key is the whole question: - **Tensor arguments** contribute `(dtype, shape)`. Shape is compared dimension by dimension, so `(32, 3)` and `(64, 3)` are different keys unless a dimension has been declared symbolic. - **Python arguments** (ints, floats, strings, bools, lists, objects) contribute their *value* — or, for arbitrary objects, their identity. Every distinct value gets its own graph, because the value was baked into the graph as a constant when the body was traced. - Nested structures are flattened, so a dict of tensors contributes the keys plus each tensor's dtype and shape. When no cached graph matches, TensorFlow traces again: it reruns the Python body on fresh symbolic placeholders and stores another graph. ## The symptoms TensorFlow itself tells you, once it has noticed: `WARNING:tensorflow: 5 out of the last 5 calls to <function f> triggered tf.function retracing. Tracing is expensive and the excessive number of tracings could be due to (1) creating @tf.function repeatedly in a loop, (2) passing tensors with different shapes, (3) passing Python objects instead of tensors.` Those three causes are exactly the ones you will be asked about. Beyond the warning, the tells are: a training step that never speeds up after warm-up, memory that grows steadily without any tensor growing, and a trace-time `print` in the body that keeps firing. ## Cause 1 — a decorator inside a loop ``` for batch in data: step = tf.function(compute) # a brand-new function object each time step(batch) ``` Each `tf.function(...)` call creates a *new* object with an *empty* cache, so nothing is ever reused. Hoist the decoration out of the loop; use `@tf.function` at definition site. ## Cause 2 — varying shapes Ragged batch sizes are the common case: the last batch of an epoch is short, or sequence length varies per batch. Each new shape is a new graph. Options: - **`input_signature`** — `@tf.function(input_signature=[tf.TensorSpec(shape=[None, 3], dtype=tf.float32)])`. A `None` dimension is symbolic: the graph is traced once with an unknown batch size and reused for every batch size. The trade-off is that the compiler knows less, so some shape-dependent optimizations are unavailable; and any op needing a concrete batch size must read it with `tf.shape(x)[0]` at run time rather than `x.shape[0]`. - **`reduce_retracing=True`** — TensorFlow observes the shapes it has seen and generalizes them for you, relaxing the specialized dimensions after a couple of variants. Less control than `input_signature`, no spec to write. - **Padding or bucketing** the inputs so only a handful of distinct shapes ever reach the function. ## Cause 3 — Python arguments ``` @tf.function def power(x, n): return x ** n for n in range(1, 100): power(tf.constant(2.0), n) # 99 graphs ``` `n` is a Python int, so it is part of the key and is folded into the graph as a constant. Passing `tf.constant(float(n))` instead makes it a tensor input: one graph, called 99 times. The rule of thumb is that anything that *varies per call* should be a tensor, and anything that is genuinely structural (a boolean flag choosing between two architectures, a layer count) is legitimately a Python argument whose few values deserve their own specialized graphs. A related subtlety: a Python `bool` argument used in an `if` is *supposed* to specialize, because the branch is baked in and both variants are cheap. Two graphs for `training=True/False` is the intended behaviour, not a bug. ## Concrete functions `f.get_concrete_function(tf.TensorSpec([None, 3], tf.float32))` returns the single `ConcreteFunction` for that signature — a graph with fixed input types. It is useful in three ways: it forces tracing at a moment you choose rather than in the middle of a latency-sensitive first request; it is the object you inspect (`cf.structured_input_signature`, `cf.graph`) when you want to know what was actually captured; and it is what export machinery serializes, which is why exporting a model means naming the input specs. ## Cost of getting it wrong Tracing runs your Python and builds and optimizes a graph — orders of magnitude more expensive than a call. Worse, every graph stays alive in the cache for as long as the function object does, so pathological retracing is also a slow memory leak. In a serving process where request shapes vary, an unbounded cache is a genuine production incident, and pinning `input_signature` is the standard defence.