Introduction

In modern TensorFlow, the usual way to build a graph from a Python function is @tf.function. You write normal TensorFlow code in a function, TensorFlow traces it into a graph, and then reuses that graph for execution. This is the current graph-building model in TensorFlow 2, and it is very different from the old manual tf.Graph plus session style.

Start with eager mode, then trace with tf.function

TensorFlow 2 runs eagerly by default, which means operations execute immediately and are easier to debug. When you decorate a function with @tf.function, TensorFlow traces the TensorFlow operations inside that function and builds a graph representation for faster repeated execution.

1import tensorflow as tf
2
3@tf.function
4def add_and_square(x, y):
5    z = x + y
6    return z * z
7
8result = add_and_square(tf.constant(2.0), tf.constant(3.0))
9print(result.numpy())

This looks like ordinary Python, but under the hood TensorFlow has the option to compile and reuse a graph.

Why build graphs inside functions at all

A graph gives TensorFlow more room to optimize execution, place operations on devices, and reduce Python overhead in repeated training or inference loops. That is why model training steps, loss functions, and custom layers often benefit from tf.function.

The practical pattern is simple:

• write clear TensorFlow code first
• confirm it works eagerly
• wrap stable hot paths in @tf.function

That order matters because debugging traced code is harder than debugging eager code.

Be aware of tracing and retracing

TensorFlow traces a function based on the TensorFlow operations it sees and, in many cases, the input signature or input shapes and dtypes. If you call the function with meaningfully different inputs, TensorFlow may retrace and build additional graphs.

1import tensorflow as tf
2
3@tf.function
4def scale(x):
5    return x * 2
6
7print(scale(tf.constant([1.0, 2.0])))
8print(scale(tf.constant([1.0, 2.0, 3.0])))

Retracing is not always wrong, but excessive retracing can hurt performance. If the function is called frequently with predictable shapes, providing an input signature can help.

@tf.function(input_signature=[tf.TensorSpec(shape=[None], dtype=tf.float32)])
def stable_scale(x):
    return x * 2

Do not create new variables on every call

One of the most common mistakes is trying to create tf.Variable objects inside a traced function every time it runs. TensorFlow expects variables to be created once and then reused.

Bad pattern:

1@tf.function
2def bad(x):
3    weight = tf.Variable(2.0)
4    return x * weight

Better pattern:

1import tensorflow as tf
2
3class Scaler(tf.Module):
4    def __init__(self):
5        self.weight = tf.Variable(2.0)
6
7    @tf.function
8    def __call__(self, x):
9        return x * self.weight
10
11scaler = Scaler()
12print(scaler(tf.constant(3.0)).numpy())

Keep variable creation outside the repeatedly traced execution path.

Use tf.Graph directly only for special cases

You can still build graphs explicitly with tf.Graph() in compatibility scenarios, but that is no longer the normal TensorFlow 2 style.

1import tensorflow as tf
2
3graph = tf.Graph()
4with graph.as_default():
5    x = tf.constant(2.0)
6    y = tf.constant(3.0)
7    z = x + y

This style makes sense mostly when dealing with older TensorFlow 1.x concepts, specialized tooling, or imported graph definitions. For ordinary new code, tf.function is the cleaner approach.

Keep Python side effects out of the critical path

Tracing works best when the function body is mostly TensorFlow operations. Python-side side effects such as appending to lists, printing for logic, or changing ordinary Python state may behave unexpectedly because the graph captures TensorFlow computation, not arbitrary Python semantics in the way beginners often imagine.

So when building graphs inside functions, think in tensors and TensorFlow ops first.

Common Pitfalls

• Reaching for manual graph construction when tf.function is the correct TensorFlow 2 tool.
• Creating new tf.Variable objects inside a traced function on repeated calls.
• Debugging a complex traced function before first validating the eager version.
• Ignoring retracing overhead when shapes or Python inputs vary frequently.
• Expecting ordinary Python side effects to behave like TensorFlow graph operations.

Summary

• In TensorFlow 2, build graphs from functions with @tf.function.
• Start in eager mode, then trace stable performance-critical code.
• Watch out for retracing when input shapes or signatures vary.
• Create variables outside the repeatedly traced function body.
• Use manual tf.Graph construction only for compatibility or specialized cases.