TensorFlow
Dataset
from_generator
pyfunc
exception

Tensorflow Dataset.from_generator fails with pyfunc exception

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

Introduction

A pyfunc exception from tf.data.Dataset.from_generator usually means the Python generator failed or returned data that does not match the declared output signature. The stack trace often looks like a TensorFlow internal failure, which can hide the real source. The fastest fix is to validate generator output contract step by step before running full training.

Why pyfunc Appears in the Error Path

from_generator bridges Python code into the tf.data pipeline through a callback boundary. If the generator raises an exception or yields invalid structure, TensorFlow reports failure at that callback layer.

Typical causes:

  • Yielded shape differs from declared TensorSpec.
  • Dtype varies across records.
  • Tuple or dict structure is inconsistent.
  • Data parsing code inside generator throws runtime error.

So the issue is often generator contract, not TensorFlow runtime corruption.

Start from a Minimal Reproducible Generator

Create a small baseline that is guaranteed correct.

python
1import numpy as np
2import tensorflow as tf
3
4
5def gen_ok():
6    for _ in range(3):
7        yield np.array([1.0, 2.0], dtype=np.float32)
8
9sig = tf.TensorSpec(shape=(2,), dtype=tf.float32)
10ds = tf.data.Dataset.from_generator(gen_ok, output_signature=sig)
11
12for x in ds:
13    print(x)

If this works, add your real preprocessing code incrementally until failure reappears.

Fix Shape Mismatch Errors

A common bug is declaring fixed shape but yielding variable-length arrays.

python
def gen_bad_shape():
    yield np.array([1.0, 2.0], dtype=np.float32)
    yield np.array([3.0], dtype=np.float32)

If length truly varies, declare flexible dimension and normalize later.

python
sig = tf.TensorSpec(shape=(None,), dtype=tf.float32)

For fixed-shape models, pad or truncate before yield.

Fix Dtype Drift Explicitly

Mixed upstream data sources often produce float64, int64, or object arrays unexpectedly. Cast at boundary.

python
1def gen_cast(rows):
2    for r in rows:
3        x = np.asarray(r["x"], dtype=np.float32)
4        y = np.asarray(r["y"], dtype=np.int32)
5        yield x, y

Do not rely on implicit conversion during pipeline execution.

Validate Structure Before Yield

When yielding multi-item records, structure must remain identical every iteration.

python
1output_signature = (
2    tf.TensorSpec(shape=(128,), dtype=tf.float32),
3    tf.TensorSpec(shape=(), dtype=tf.int32),
4)
5
6
7def gen_records(records):
8    for rec in records:
9        x = np.asarray(rec["x"], dtype=np.float32)
10        y = np.int32(rec["y"])
11
12        if x.shape != (128,):
13            raise ValueError(f"unexpected shape {x.shape}")
14
15        yield x, y

Early checks produce actionable errors instead of opaque callback failures.

Debug Outside tf.data First

Before building dataset, iterate generator directly.

python
g = gen_ok()
print(next(g))

Then test dataset with a small slice:

python
for item in ds.take(1):
    print(item)

This two-stage debug approach separates Python data issues from pipeline composition issues.

Production Hardening

For long training jobs:

  • Log record ids before parsing.
  • Catch and re-raise with contextual metadata.
  • Decide policy for malformed records, either fail fast or skip with metrics.

Silent skips without metrics can hide data quality regressions and degrade model performance.

Performance Note

from_generator is convenient but Python-bound. Heavy preprocessing inside generator can become throughput bottleneck. Move expensive transforms into native tf.data ops when possible and add prefetch.

python
ds = ds.prefetch(tf.data.AUTOTUNE)

Correctness first, then throughput optimization.

Common Pitfalls

  • Declaring output signature that does not match actual yields. Fix by aligning shape, dtype, and structure exactly.
  • Allowing one record to emit a different dtype. Fix with explicit casting before yield.
  • Debugging only TensorFlow stack trace. Fix by testing generator directly in plain Python.
  • Ignoring malformed record context. Fix by logging record identifiers and validation failures.
  • Overloading generator with expensive transforms. Fix by shifting heavy work into tf.data operations.

Summary

  • 'pyfunc errors in from_generator usually indicate generator contract violations.'
  • Validate shape, dtype, and structure before data enters the dataset pipeline.
  • Build from a minimal working generator and add complexity gradually.
  • Use early validation and contextual logging for fast diagnosis.
  • Keep generator lean and optimize pipeline throughput after correctness is stable.

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

All Rights Reserved.