Introduction

tf.py_function is useful when you need to bridge Python-only preprocessing into a TensorFlow pipeline, but it comes with a sharp edge: static shape information is often lost. That becomes a problem later when you batch a dataset, build a Keras model, or rely on graph tracing that expects known ranks and dimensions. The fix is usually simple, but it has to happen immediately after the tf.py_function call.

Why Shape Information Disappears

TensorFlow can infer shapes through normal TensorFlow ops because it understands those ops symbolically. tf.py_function is different: it calls arbitrary Python code, so TensorFlow cannot reason about the returned tensor shape in advance.

1import tensorflow as tf
2import numpy as np
3
4def py_preprocess(x):
5    arr = x.numpy().astype(np.float32) / 255.0
6    return arr
7
8def wrapped(x):
9    y = tf.py_function(py_preprocess, [x], Tout=tf.float32)
10    return y

After this call, y often has an unknown static shape even if the runtime value is perfectly valid.

Restore The Shape Explicitly

If you know the output shape, tell TensorFlow directly.

1def wrapped_with_shape(x):
2    y = tf.py_function(py_preprocess, [x], Tout=tf.float32)
3    y.set_shape(x.shape)
4    return y

For image pipelines with a fixed size, you can set it explicitly.

1def wrapped_image(x):
2    y = tf.py_function(py_preprocess, [x], Tout=tf.float32)
3    y.set_shape([224, 224, 3])
4    return y

This restores the static metadata that downstream TensorFlow components need.

Use tf.ensure_shape When You Want Validation Too

set_shape annotates the tensor. tf.ensure_shape both annotates and checks that the runtime shape matches the expectation.

1def wrapped_checked(x):
2    y = tf.py_function(py_preprocess, [x], Tout=tf.float32)
3    y = tf.ensure_shape(y, x.shape)
4    return y

This is useful during development because shape mismatches fail close to the source of the problem instead of surfacing later inside a model call.

A tf.data Example

Shape loss often appears most clearly inside a dataset pipeline.

1images = tf.random.uniform((10, 224, 224, 3), maxval=255, dtype=tf.int32)
2ds = tf.data.Dataset.from_tensor_slices(images)
3
4def map_fn(x):
5    x = tf.cast(x, tf.uint8)
6    y = tf.py_function(py_preprocess, [x], Tout=tf.float32)
7    y.set_shape([224, 224, 3])
8    return y
9
10ds = ds.map(map_fn).batch(4)
11
12for batch in ds.take(1):
13    print(batch.shape)

Without the explicit shape restoration, batching and model input checks can fail because TensorFlow sees an unknown shape where the pipeline actually expects a fixed image tensor.

Prefer Native TensorFlow Ops When Possible

tf.py_function should be a bridge, not your first choice. Native TensorFlow ops are better because they preserve shape information, work more cleanly with graph tracing, and are more portable for saved models and serving.

If the preprocessing can be expressed with tf.image, tf.strings, tf.cast, or other TensorFlow ops, use those instead. Reserve tf.py_function for genuinely Python-only logic or legacy code you have not replaced yet.

Watch Dtype And Rank At The Same Time

Shape problems are often accompanied by dtype problems. If your callback returns a NumPy array with the wrong dtype or rank, a later layer may fail in ways that look unrelated to the original preprocessing step.

A good debugging pattern is to inspect both:

1def map_fn(x):
2    y = tf.py_function(py_preprocess, [x], Tout=tf.float32)
3    y.set_shape([224, 224, 3])
4    tf.print("shape:", tf.shape(y), "dtype:", y.dtype)
5    return y

That makes it easier to catch the real issue before a long training run hides it inside a deeper stack trace.

Export And Serving Concerns

Python callbacks are also a portability concern. A pipeline that depends on arbitrary Python code may not export or serve cleanly in all TensorFlow environments. If the model needs to be saved and deployed broadly, replacing tf.py_function with native TensorFlow preprocessing is usually worth the effort even when the quick local fix is just set_shape.

Common Pitfalls

• Assuming tf.py_function preserves static shape information automatically.
• Fixing only the dtype and forgetting to restore the shape metadata.
• Setting the wrong shape manually and hiding a deeper preprocessing bug.
• Using tf.py_function for logic that could have been written with native TensorFlow ops.
• Discovering the problem only at model-training time instead of validating shape immediately after preprocessing.

Summary

• 'tf.py_function often drops static shape information because TensorFlow cannot infer arbitrary Python behavior.'
• Restore the expected shape right after the call with set_shape or tf.ensure_shape.
• Validate dtype and rank early, especially inside tf.data pipelines.
• Prefer native TensorFlow ops whenever possible for better tracing and portability.
• Treat tf.py_function as a compatibility bridge, not the default preprocessing strategy.