Practice Now
TensorFlow
py_func
unknown shape
machine learning
Python

Tensorflow Py_func returns unknown shape

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design
TensorFlow & KerasMachine LearningPython

Introduction

tf.py_function wraps arbitrary Python code into a TensorFlow operation, but the resulting tensors always have unknown (None) shapes because TensorFlow cannot infer shapes from opaque Python code. The fix is to call tensor.set_shape() on each output tensor after calling tf.py_function to manually restore the shape information. Without this, downstream operations that require known shapes (like Dense layers, batch(), or reshape()) fail.

The Problem

python
1import tensorflow as tf
2import numpy as np
3
4def my_preprocess(image, label):
5    # Custom Python preprocessing
6    image = image.numpy()  # Convert to numpy
7    image = image / 255.0
8    return image.astype(np.float32), label.numpy()
9
10# Wrap with tf.py_function
11dataset = tf.data.Dataset.from_tensor_slices(
12    (tf.random.uniform([100, 28, 28, 1], maxval=255, dtype=tf.int32),
13     tf.random.uniform([100], maxval=10, dtype=tf.int32))
14)
15
16dataset = dataset.map(lambda img, lbl: tf.py_function(
17    my_preprocess, [img, lbl], [tf.float32, tf.int32]
18))
19
20# Check shapes — they are unknown!
21for img, lbl in dataset.take(1):
22    print(img.shape)  # (None, None, None, None) — unknown!
23    print(lbl.shape)  # (None,) — unknown!

The Fix: set_shape()

python
1def my_preprocess(image, label):
2    image = image.numpy().astype(np.float32) / 255.0
3    return image, label.numpy()
4
5def preprocess_wrapper(image, label):
6    image, label = tf.py_function(
7        my_preprocess, [image, label], [tf.float32, tf.int32]
8    )
9    # Manually restore shape information
10    image.set_shape([28, 28, 1])
11    label.set_shape([])
12    return image, label
13
14dataset = dataset.map(preprocess_wrapper)
15
16# Now shapes are known
17for img, lbl in dataset.take(1):
18    print(img.shape)  # (28, 28, 1)
19    print(lbl.shape)  # ()

Why Shapes Are Unknown

TensorFlow builds a computation graph before execution. Regular TF ops declare their output shapes based on input shapes (e.g., Conv2D knows its output shape from input shape, kernel size, and padding). tf.py_function runs arbitrary Python code that TensorFlow cannot analyze, so it conservatively sets all output shapes to unknown.

python
1# Regular TF op — shapes are known
2x = tf.random.normal([32, 28, 28, 1])
3y = tf.keras.layers.Conv2D(16, 3)(x)
4print(y.shape)  # (32, 26, 26, 16) — TF infers this
5
6# py_function — shapes are unknown
7def identity(x):
8    return x.numpy()
9
10result = tf.py_function(identity, [x], tf.float32)
11print(result.shape)  # (None, None, None, None) — unknown

Using with tf.data Pipelines

python
1def load_and_augment(filepath, label):
2    """Custom augmentation in pure Python."""
3    def _process(filepath, label):
4        path = filepath.numpy().decode("utf-8")
5        img = np.random.rand(224, 224, 3).astype(np.float32)  # Placeholder
6        return img, label.numpy()
7
8    image, label = tf.py_function(
9        _process, [filepath, label], [tf.float32, tf.int32]
10    )
11    image.set_shape([224, 224, 3])
12    label.set_shape([])
13    return image, label
14
15# Create dataset
16files = tf.constant(["img1.jpg", "img2.jpg", "img3.jpg"])
17labels = tf.constant([0, 1, 2])
18dataset = tf.data.Dataset.from_tensor_slices((files, labels))
19
20dataset = dataset.map(load_and_augment, num_parallel_calls=tf.data.AUTOTUNE)
21dataset = dataset.batch(2)  # Works because shapes are known
22
23for images, labels in dataset.take(1):
24    print(images.shape)  # (2, 224, 224, 3)
25    print(labels.shape)  # (2,)

tf.py_function vs tf.numpy_function

python
1# tf.py_function — inputs are tf.Tensor (call .numpy() to get values)
2def with_py_function(x):
3    return x.numpy() * 2  # Must call .numpy()
4
5# tf.numpy_function — inputs are already numpy arrays
6def with_numpy_function(x):
7    return x * 2  # Already a numpy array
8
9result1 = tf.py_function(with_py_function, [tf.constant([1, 2])], tf.int32)
10result2 = tf.numpy_function(with_numpy_function, [tf.constant([1, 2])], tf.int32)
11
12# Both have unknown shapes — both need set_shape()
13result1.set_shape([2])
14result2.set_shape([2])

ensure_shape() as an Alternative

python
1# ensure_shape raises an error at runtime if shape doesn't match
2image = tf.ensure_shape(image, [28, 28, 1])
3label = tf.ensure_shape(label, [])
4
5# Difference:
6# set_shape() — compile-time assertion, no runtime check
7# ensure_shape() — adds a runtime shape check op

Use ensure_shape() when you want to catch shape mismatches at runtime (debugging), and set_shape() when you are confident about the shape.

Common Pitfalls

  • Forgetting to call set_shape() after tf.py_function: Without it, all downstream ops see unknown shapes. model.fit() with a tf.data pipeline will fail if the model expects known input shapes. Always set shapes immediately after tf.py_function.
  • Setting the wrong shape in set_shape(): If you set image.set_shape([28, 28, 3]) but the function actually returns (28, 28, 1), TensorFlow does not catch this at graph-build time. The error only appears at runtime as a shape mismatch. Use tf.ensure_shape() during development to catch mistakes.
  • Using tf.py_function inside a model's call() method: tf.py_function runs Python code eagerly and cannot be traced by tf.function or exported with SavedModel. Use it only in tf.data pipelines for preprocessing, not inside model layers.
  • Not calling .numpy() on inputs inside tf.py_function: Inside a tf.py_function, inputs are tf.Tensor objects in eager mode. Call .numpy() to get numpy arrays before passing to libraries like OpenCV or scipy. Forgetting this causes type errors.
  • Performance degradation from Python GIL: tf.py_function runs Python code and holds the GIL, preventing true parallelism. For data pipelines, use num_parallel_calls=tf.data.AUTOTUNE and prefetch() to overlap Python execution with GPU training. For performance-critical code, rewrite in pure TensorFlow ops.

Summary

  • tf.py_function outputs always have unknown shapes — call set_shape() on each output tensor
  • Use tf.ensure_shape() during development for runtime shape validation
  • Only use tf.py_function in tf.data pipelines, not inside model layers
  • Call .numpy() on tensor inputs inside the wrapped function to get numpy arrays
  • Use num_parallel_calls and prefetch() to mitigate the GIL bottleneck
Related reading
Free course
Beginner
7 lessons
2 hours
Tackling System Design Interview Problems

A short course that equips you with the skills to approach system design interviews methodically.

Start the free course
Track what you have practised

A free account saves your progress, solutions and study plan across every problem on Codemia.

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design
Codemia
© 2026 Codemia
Resources
Legal
Social
All Rights Reserved.