Common tf.function Mistakes: Python Side Effects and Retracing

The Story — The Recipe Card With Today's Date

This lesson covers the five mistakes that almost every TensorFlow learner makes with tf.function. Each one has the same root cause: mixing up trace time (Python runs) and run time (only the graph runs).

Mistake 1 — Python Side Effects

A side effect is anything a function does besides returning a value: printing, adding to a list, changing a global counter.

import tensorflow as tf

calls = 0
results = []

@tf.function
def step(x):
    global calls
    calls += 1                 # Python side effect
    results.append(x * 2)      # Python side effect
    return x * 2

for i in range(3):
    step(tf.constant(i))

print("calls counted :", calls)
print("results list  :", results)

We called the function 3 times, but the counter says 1. The list holds one symbolic tensor, not three numbers. Python ran only while tracing.

The Fix: Keep State in TensorFlow

calls = tf.Variable(0)

@tf.function
def step(x):
    calls.assign_add(1)                           # a graph op: runs every call
    return x * 2

outputs = [int(step(tf.constant(i))) for i in range(3)]   # collect OUTSIDE
print("calls counted :", calls.numpy())
print("outputs       :", outputs)

To collect values inside a graph loop, use tf.TensorArray:

@tf.function
def squares(n):
    ta = tf.TensorArray(tf.int32, size=n)
    for i in tf.range(n):
        ta = ta.write(i, i * i)
    return ta.stack()

print(squares(tf.constant(5)).numpy())

Mistake 2 — NumPy Randomness and the Clock Are Frozen

import numpy as np

@tf.function
def add_noise(x):
    return x + np.random.randn()      # NumPy runs once -> a constant

print("NumPy noise:", [round(float(add_noise(tf.constant(0.0))), 4) for _ in range(3)])

@tf.function
def add_noise_tf(x):
    return x + tf.random.normal([])   # a TF op -> new value each call

print("TF noise   :", [round(float(add_noise_tf(tf.constant(0.0))), 4) for _ in range(3)])

Mistake 3 — Creating Variables Inside the Function

@tf.function
def bad_layer(x):
    w = tf.Variable(2.0)       # a new variable on every trace
    return w * x

try:
    bad_layer(tf.constant(3.0))
except ValueError as e:
    print("ValueError:", str(e).strip().splitlines()[-1].split("ValueError: ")[-1][:95], "...")

Create variables outside the function, or only once — the same pattern Keras layers use in build().

class Scale(tf.Module):
    def __init__(self):
        self.w = None

    @tf.function
    def __call__(self, x):
        if self.w is None:                 # create only on the first trace
            self.w = tf.Variable(2.0)
        return self.w * x

s = Scale()
print(s(tf.constant(3.0)).numpy(), s(tf.constant(4.0)).numpy())

Mistake 4 — Calling .numpy() Inside

@tf.function
def to_list(x):
    return x.numpy()           # no values exist during tracing

try:
    to_list(tf.constant([1, 2]))
except AttributeError as e:
    print("AttributeError:", str(e).strip().splitlines()[-1].split("AttributeError: ")[-1])

Inside a tf.function, tensors are symbolic. Keep the maths in TF ops, and call .numpy() on the result, outside the function.

Mistake 5 — Retracing Again and Again

@tf.function
def double(x):
    return x * 2

for i in range(5):
    double(i)                  # Python ints: each new value is a new signature
print("Python ints  -> traces:", double.experimental_get_tracing_count())

@tf.function
def double_t(x):
    return x * 2

for i in range(5):
    double_t(tf.constant(i))   # tensors: same shape and dtype every time
print("Tensors      -> traces:", double_t.experimental_get_tracing_count())

After 5 retraces, TensorFlow prints a warning like this in your log:

Try It — Bug Hunter

Try It — Retrace Counter

Golden Rules