Introduction

tf.function and tf.py_function both sit at the boundary between Python and TensorFlow execution, but they solve very different problems. tf.function turns TensorFlow-heavy Python code into a TensorFlow graph for tracing and optimization. tf.py_function does the opposite kind of bridge: it lets arbitrary Python code appear inside TensorFlow execution when no pure TensorFlow op is available.

tf.function: trace TensorFlow code into a graph

The purpose of tf.function is to take Python code that uses TensorFlow ops and trace it into a callable graph.

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), tf.constant(3))
9print(result)

This is useful because TensorFlow can optimize the traced computation, reuse it, and run it in contexts where graph execution matters.

The important point is that tf.function works best when the body is already written in TensorFlow terms. The more Python-side control flow and Python objects dominate the function, the less benefit you get and the more surprising the tracing behavior can become.

tf.py_function: inject raw Python into TensorFlow execution

tf.py_function exists for a different reason. It lets you run a regular Python function as part of TensorFlow execution, usually when you need logic that is not naturally expressed as TensorFlow ops.

1import tensorflow as tf
2
3def double_in_python(x):
4    return x * 2
5
6tensor = tf.constant([1, 2, 3], dtype=tf.int32)
7result = tf.py_function(func=double_in_python, inp=[tensor], Tout=tf.int32)
8print(result)

This is handy for quick experimentation, legacy Python code, or custom data processing that is difficult to rewrite in TensorFlow. But it comes with tradeoffs: the wrapped code runs in Python, is harder for TensorFlow to optimize, and is not a substitute for a proper TensorFlow op when performance or portability matters.

Their purposes are almost opposite

A simple way to remember the distinction is:

• 'tf.function tries to move your code deeper into the TensorFlow graph world.'
• 'tf.py_function lets Python leak back into TensorFlow execution when necessary.'

That is why they are not interchangeable. One is about graph tracing and optimization. The other is about compatibility with arbitrary Python logic.

Typical use cases

Use tf.function when:

• Your function is mainly TensorFlow ops
• You want graph execution and tracing benefits
• You are building training, inference, or reusable computation code

Use tf.py_function when:

• You must call ordinary Python code
• You are prototyping quickly
• You need a stopgap for custom logic inside a TensorFlow pipeline

For example, in a dataset pipeline:

1import tensorflow as tf
2
3def py_map(x):
4    return x + 5
5
6dataset = tf.data.Dataset.from_tensor_slices([1, 2, 3])
7dataset = dataset.map(lambda x: tf.py_function(py_map, [x], Tout=tf.int32))
8
9for item in dataset:
10    print(item)

This works, but it is usually a fallback, not the ideal long-term implementation.

Why tf.py_function is not a performance tool

Because tf.py_function runs Python code, it keeps execution tied to the Python runtime. That limits graph portability and can become a bottleneck. If the logic can be written using TensorFlow ops, that is usually the better direction.

By contrast, tf.function is specifically about helping TensorFlow capture the computation rather than falling back to arbitrary Python execution.

Common Pitfalls

The biggest mistake is treating tf.py_function as if it were a faster or more graph-friendly version of a normal function. It is usually less optimizable, not more.

Another issue is wrapping Python-heavy code with tf.function and expecting all of it to behave like plain eager Python. Tracing can introduce subtleties around control flow, side effects, and retracing.

Developers also use tf.py_function in data pipelines without realizing that it can reduce portability and complicate performance tuning.

Finally, do not confuse "both accept functions" with "both solve the same problem." The similarity in naming hides a very different purpose.

Summary

• 'tf.function traces TensorFlow-based code into a TensorFlow graph.'
• 'tf.py_function wraps arbitrary Python code so it can run during TensorFlow execution.'
• 'tf.function is for graph execution and optimization.'
• 'tf.py_function is for bridging to Python when no convenient TensorFlow op exists.'
• If pure TensorFlow ops can express the logic, prefer that over tf.py_function.