How to compile a function with tf.function in TensorFlow

TensorFlow code often starts in eager mode, where Python runs each operation immediately and makes debugging straightforward. Wrapping the repeated tensor-only part with tf.function lets TensorFlow trace that Python function into a graph-backed callable, which reduces Python overhead for stable inference helpers, preprocessing functions, and custom training steps.

tf.function creates a polymorphic callable, which means the decorated Python function can cache concrete graphs for the tensor shapes, dtypes, and Python argument values it sees. An input_signature pins the accepted tensor contract, and get_concrete_function() exposes the concrete graph for that contract.

Graph tracing with tf.function is separate from XLA compilation. Add jit_compile=True only when the function's operations support XLA and the extra compiler pass is the actual target. Compile after the eager version returns the expected tensors, keep variable and layer creation outside repeated calls, and pass changing numeric controls as tensors when the call should stay on one graph.

Steps to compile a function with tf.function in TensorFlow:

Save a scoring demo as

tf_function_compile_demo.py

tf_function_compile_demo.py

import os
 
os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2"
 
import tensorflow as tf
 
tf.get_logger().setLevel("ERROR")
 
weights = tf.constant([0.7, -0.2, 0.5], dtype=tf.float32)
bias = tf.constant(0.1, dtype=tf.float32)
 
 
@tf.function(
    input_signature=[
        tf.TensorSpec(shape=(None, 3), dtype=tf.float32, name="features")
    ]
)
def score_batch(features):
    return tf.reduce_sum(features * weights, axis=1) + bias
 
 
compiled_score = score_batch.get_concrete_function()
 
first_batch = tf.constant(
    [
        [1.0, 0.2, 0.5],
        [0.3, 0.9, 1.2],
    ],
    dtype=tf.float32,
)
second_batch = tf.constant([[0.8, 0.1, 0.4]], dtype=tf.float32)
 
first_scores = score_batch(first_batch)
second_scores = score_batch(second_batch)
same_graph = score_batch.get_concrete_function() is compiled_score
 
try:
    score_batch(tf.constant([[1.0, 2.0]], dtype=tf.float32))
except TypeError:
    bad_shape_error = "TypeError"
else:
    bad_shape_error = "none"
 
graph_ops = ", ".join(op.name for op in compiled_score.graph.get_operations())
 
print(f"TensorFlow {tf.__version__}")
print(f"Input signature: {compiled_score.structured_input_signature[0][0]}")
print(f"Output shape: {tuple(compiled_score.output_shapes.as_list())}")
print(f"Graph operations: {graph_ops}")
print(
    "First batch scores: "
    f"{[round(float(value), 2) for value in first_scores.numpy()]}"
)
print(
    "Second batch scores: "
    f"{[round(float(value), 2) for value in second_scores.numpy()]}"
)
print(f"Reused compiled graph: {same_graph}")
print(f"Bad shape error: {bad_shape_error}")

The None batch dimension accepts different row counts while keeping each row fixed at three float32 features.

Run the demo and confirm that the changed batch size reuses the compiled graph.

$ python3 tf_function_compile_demo.py
TensorFlow 2.21.0
Input signature: TensorSpec(shape=(None, 3), dtype=tf.float32, name='features')
Output shape: (None,)
Graph operations: features, mul/y, mul, Sum/reduction_indices, Sum, add/y, add, Identity
First batch scores: [1.01, 0.73]
Second batch scores: [0.84]
Reused compiled graph: True
Bad shape error: TypeError

Reused compiled graph: True means the one-row and two-row calls matched the same TensorSpec. Bad shape error: TypeError means the two-column tensor was rejected by the signature before graph execution.

Apply the same signature pattern to the project function that runs repeatedly.
```
@tf.function(
    input_signature=[
        tf.TensorSpec(shape=(None, 128), dtype=tf.float32, name="features")
    ]
)
def predict_batch(features):
    return model(features, training=False)
```
Use None only for dimensions that may vary between calls. Fixed dimensions catch the wrong feature width before the compiled function runs.

Related: How to run a custom training loop in TensorFlow
Materialize the concrete function when the signature needs to be inspected or exported.
```
compiled_predict = predict_batch.get_concrete_function()
print(compiled_predict.structured_input_signature)
```
get_concrete_function() returns the graph-backed callable for the matching signature. Export and serving workflows can use that signature as the model contract.

Related: How to export serving signatures in TensorFlow

Pass frequently changing numeric options as tensors when the call should stay on one graph.

@tf.function(
    input_signature=[
        tf.TensorSpec(shape=(None,), dtype=tf.float32, name="scores"),
        tf.TensorSpec(shape=(), dtype=tf.float32, name="temperature"),
    ]
)
def scale_scores(scores, temperature):
    return scores / temperature
 
 
scores = tf.constant([0.35, 0.80], dtype=tf.float32)
temperature = tf.constant(0.7, dtype=tf.float32)
scaled = scale_scores(scores, temperature)

Python arguments are treated as compile-time values by tf.function. New Python values can create extra concrete graphs even when the tensor math is unchanged.

Create variables, layers, and optimizers before repeated compiled calls begin.

weights = tf.Variable(tf.ones((3,), dtype=tf.float32))
 
 
@tf.function(
    input_signature=[
        tf.TensorSpec(shape=(None, 3), dtype=tf.float32, name="features")
    ]
)
def score(features):
    return tf.linalg.matvec(features, weights)

tf.function allows new tf.Variable objects only on the first call, so creating variables, layers, or optimizer state inside an already-cached function can fail on later calls.

Remove the demo script after the project function has its own compiled signature.
```
$ rm tf_function_compile_demo.py
```