Introduction

When you use tf.py_function, tf.numpy_function, or the older tf.py_func, TensorFlow executes Python code but loses static shape information about the returned tensor. The result often has unknown rank or unknown dimensions, which then breaks later layers, batching, or tracing. The fix is usually to restore the shape explicitly after the Python callback.

Why TensorFlow Loses the Shape

TensorFlow can infer shapes for built-in ops because their semantics are known to the graph runtime. A Python callback is opaque: TensorFlow knows the output dtype you declared, but not the actual tensor shape.

Example:

1import tensorflow as tf
2
3
4def double_numpy(x):
5    return x * 2
6
7
8x = tf.constant([1, 2, 3], dtype=tf.int32)
9y = tf.py_function(double_numpy, [x], Tout=tf.int32)
10
11print(y.shape)

y.shape is usually unknown or only partially known even though the runtime value is fine.

That missing metadata becomes a real problem when a later operation expects a specific rank or dimension.

Restore Shape with set_shape or ensure_shape

If you know the output shape, tell TensorFlow directly.

1import tensorflow as tf
2
3
4def double_numpy(x):
5    return x * 2
6
7
8x = tf.constant([1, 2, 3], dtype=tf.int32)
9y = tf.py_function(double_numpy, [x], Tout=tf.int32)
10y.set_shape([3])
11
12print(y.shape)

In newer code, tf.ensure_shape is often clearer because it both annotates and validates:

1import tensorflow as tf
2
3
4def double_numpy(x):
5    return x * 2
6
7
8x = tf.constant([1, 2, 3], dtype=tf.int32)
9y = tf.py_function(double_numpy, [x], Tout=tf.int32)
10y = tf.ensure_shape(y, [3])
11
12print(y.shape)

If the runtime output does not match that declared shape, TensorFlow raises an error instead of letting the mismatch propagate silently.

This Is Common in tf.data Pipelines

The issue often appears inside Dataset.map, especially when a Python function is used for custom parsing.

1import tensorflow as tf
2
3
4def parse_py(x):
5    return x + 1
6
7
8def wrapped(x):
9    y = tf.py_function(parse_py, [x], Tout=tf.int32)
10    y.set_shape([])
11    return y
12
13
14ds = tf.data.Dataset.from_tensor_slices([1, 2, 3])
15ds = ds.map(wrapped)
16
17for item in ds:
18    print(item.numpy())

Without set_shape([]), later batching or model input validation may complain because TensorFlow only sees an unknown scalar-like object instead of a declared scalar tensor.

Prefer Native TensorFlow Ops When Possible

py_function is useful, but it should usually be a last resort. Native TensorFlow ops preserve:

• shape inference
• graph optimization
• portability
• better tracing behavior

For example, if the transformation can be expressed with TensorFlow directly:

1import tensorflow as tf
2
3x = tf.constant([1, 2, 3], dtype=tf.int32)
4y = x * 2
5print(y.shape)

This is much better than using py_function for something TensorFlow already supports.

Be Careful with Rank Versus Dimension

Unknown rank and unknown dimension are different problems:

• unknown rank means TensorFlow does not even know how many axes the tensor has
• unknown dimension means the number of axes is known, but one or more sizes are not

For example:

• shape [None, 128] has known rank two but an unknown first dimension
• shape TensorShape(None) may indicate rank is unknown

That distinction matters because some layers only need known rank, while others need specific dimensions too.

Modern Replacement for tf.py_func

In TensorFlow 2, prefer:

• 'tf.py_function for eager-style tensor input and output'
• 'tf.numpy_function when you specifically want NumPy arrays inside the callback'

The shape problem remains the same. Neither op can infer the full output shape automatically, so explicit shape restoration is still part of the solution.

Common Pitfalls

• Expecting py_function to preserve the input shape automatically.
• Using Python callbacks for work that could be done with native TensorFlow ops.
• Forgetting to restore shape inside tf.data pipelines.
• Confusing unknown rank with partially known dimensions.
• Returning inconsistent shapes from the Python function and then annotating the wrong one.

Summary

• TensorFlow cannot infer full shape information from py_function or the older py_func.
• The output dtype is known, but the shape often is not.
• Use set_shape or tf.ensure_shape immediately after the Python callback.
• Prefer native TensorFlow ops whenever you can.
• In tf.data and model pipelines, explicit shape restoration is usually the difference between working code and confusing downstream errors.