How do you send arguments to a generator function using tf.data.Dataset.from_generator?

TensorFlow

generator functions

tf.data.Dataset

from_generator

Python

How do you send arguments to a generator function using tf.data.Dataset.from_generator?

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 & Keras Python

Introduction

tf.data.Dataset.from_generator() creates a TensorFlow dataset from a Python generator function. To pass arguments to the generator, use the args parameter, which accepts a tuple of values that get forwarded to the generator function when the dataset iterates. Each argument in args is converted to a tf.Tensor internally, so the generator receives tensor-compatible values. This is essential for creating parameterized data pipelines — for example, passing file paths, batch sizes, or split names to control what data the generator yields.

Basic Usage Without Arguments

python

1import tensorflow as tf
2
3def simple_generator():
4    for i in range(5):
5        yield i, i * 2
6
7dataset = tf.data.Dataset.from_generator(
8    simple_generator,
9    output_signature=(
10        tf.TensorSpec(shape=(), dtype=tf.int32),
11        tf.TensorSpec(shape=(), dtype=tf.int32),
12    )
13)
14
15for x, y in dataset:
16    print(f"x={x.numpy()}, y={y.numpy()}")
17# x=0, y=0
18# x=1, y=2
19# x=2, y=4
20# x=3, y=6
21# x=4, y=8

Passing Arguments with `args`

The args parameter passes a tuple of arguments to the generator function:

python

1import tensorflow as tf
2import numpy as np
3
4def data_generator(start, end, step):
5    # Arguments arrive as numpy values (auto-converted from tensors)
6    start = start.numpy() if hasattr(start, 'numpy') else start
7    end = end.numpy() if hasattr(end, 'numpy') else end
8    step = step.numpy() if hasattr(step, 'numpy') else step
9
10    for i in range(start, end, step):
11        yield float(i)
12
13dataset = tf.data.Dataset.from_generator(
14    data_generator,
15    args=(0, 100, 5),  # Passed to data_generator(start, end, step)
16    output_signature=tf.TensorSpec(shape=(), dtype=tf.float32)
17)
18
19print(list(dataset.take(5).as_numpy_iterator()))
20# [0.0, 5.0, 10.0, 15.0, 20.0]

Passing String Arguments

Strings require explicit dtype specification because TensorFlow defaults to tf.int32:

python

1import tensorflow as tf
2
3def file_reader(file_path, label):
4    path = file_path.numpy().decode('utf-8') if hasattr(file_path, 'numpy') else file_path
5    lab = label.numpy() if hasattr(label, 'numpy') else label
6
7    # Simulate reading lines from a file
8    for i in range(3):
9        yield f"line {i} from {path}", lab
10
11dataset = tf.data.Dataset.from_generator(
12    file_reader,
13    args=("data/train.csv", 1),
14    output_signature=(
15        tf.TensorSpec(shape=(), dtype=tf.string),
16        tf.TensorSpec(shape=(), dtype=tf.int32),
17    )
18)
19
20for text, label in dataset:
21    print(text.numpy().decode(), label.numpy())

Passing Arrays and Complex Types

python

1import tensorflow as tf
2import numpy as np
3
4def batch_generator(data, labels, batch_size):
5    data = data.numpy() if hasattr(data, 'numpy') else data
6    labels = labels.numpy() if hasattr(labels, 'numpy') else labels
7    batch_size = batch_size.numpy() if hasattr(batch_size, 'numpy') else batch_size
8
9    num_samples = len(data)
10    for start in range(0, num_samples, batch_size):
11        end = min(start + batch_size, num_samples)
12        yield data[start:end], labels[start:end]
13
14X = np.random.randn(100, 10).astype(np.float32)
15y = np.random.randint(0, 3, size=100).astype(np.int32)
16
17dataset = tf.data.Dataset.from_generator(
18    batch_generator,
19    args=(X, y, 32),
20    output_signature=(
21        tf.TensorSpec(shape=(None, 10), dtype=tf.float32),
22        tf.TensorSpec(shape=(None,), dtype=tf.int32),
23    )
24)
25
26for batch_x, batch_y in dataset:
27    print(f"Batch shape: {batch_x.shape}, Labels shape: {batch_y.shape}")

Using Lambda or functools.partial

For generators that need arguments not compatible with args, use a wrapper:

python

1import tensorflow as tf
2from functools import partial
3
4def image_generator(directory, augment=False, target_size=(224, 224)):
5    # These are regular Python objects, not tensors
6    import os
7    for filename in os.listdir(directory):
8        # Simulate loading and augmenting images
9        image = tf.random.uniform(target_size + (3,))
10        if augment:
11            image = tf.image.random_flip_left_right(image)
12        yield image
13
14# Using lambda (wraps the generator with no args)
15dataset = tf.data.Dataset.from_generator(
16    lambda: image_generator("data/train", augment=True, target_size=(128, 128)),
17    output_signature=tf.TensorSpec(shape=(128, 128, 3), dtype=tf.float32)
18)
19
20# Using functools.partial
21gen_fn = partial(image_generator, "data/train", augment=True, target_size=(128, 128))
22dataset = tf.data.Dataset.from_generator(
23    gen_fn,
24    output_signature=tf.TensorSpec(shape=(128, 128, 3), dtype=tf.float32)
25)

The lambda/partial approach is preferred when arguments are Python objects (dicts, lists of strings, custom objects) that cannot be converted to tensors.

output_signature vs output_types/output_shapes

python

1# Modern API (TF 2.3+) — use output_signature
2dataset = tf.data.Dataset.from_generator(
3    my_gen,
4    args=(42,),
5    output_signature=(
6        tf.TensorSpec(shape=(None, 3), dtype=tf.float32),
7        tf.TensorSpec(shape=(), dtype=tf.int32),
8    )
9)
10
11# Legacy API — output_types and output_shapes (deprecated)
12dataset = tf.data.Dataset.from_generator(
13    my_gen,
14    args=(42,),
15    output_types=(tf.float32, tf.int32),
16    output_shapes=((None, 3), ())
17)

output_signature is preferred because it combines type and shape information into a single, clear declaration.

Common Pitfalls

Forgetting to extract .numpy() from tensor arguments: Inside the generator, arguments arrive as eager tensors (or numpy values depending on TF version). Calling .numpy() converts them to Python/NumPy types for use with standard Python operations like range(), string formatting, or indexing.
Passing non-tensor-compatible objects via args: The args tuple only accepts values that TensorFlow can convert to tensors (numbers, strings, NumPy arrays). For Python objects like dicts, lists of strings, or class instances, use a lambda or functools.partial wrapper instead.
Mismatched output_signature shapes: If your generator yields tensors with shapes that do not match output_signature, TensorFlow raises a runtime error. Use None for variable-length dimensions (e.g., shape=(None, 3) for variable-length sequences).
Generator state not resetting between epochs: from_generator creates a new generator instance each epoch by calling the generator function again. If your generator relies on external mutable state (like a file pointer), ensure it resets properly when the function is called again.
Using args with a class method: When passing a bound method as the generator, self is already captured. Additional arguments should still go in args, but ensure the method signature matches — args does not include self.

Summary

Use the args parameter to pass a tuple of arguments to the generator function in from_generator()
Arguments in args are converted to tensors — use .numpy() inside the generator to access raw values
For non-tensor-compatible arguments (dicts, custom objects), wrap the generator with a lambda or functools.partial
Always use output_signature (TF 2.3+) instead of the deprecated output_types/output_shapes parameters
The generator function is called fresh each epoch, so ensure stateful generators reset properly

Tackling System Design Interview Problems

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

Start the free course

Basis Lab

No course covers this one yet

Basis Lab writes one for you from a sentence about what you want to be able to do, then teaches it and checks you understood. Your first course is free.

Build my course

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

How do you send arguments to a generator function using tf.data.Dataset.from_generator?

ML System Design practice on Codemia

Introduction

Basic Usage Without Arguments

Passing Arguments with args

Passing String Arguments

Passing Arrays and Complex Types

Using Lambda or functools.partial

output_signature vs output_types/output_shapes

Common Pitfalls

Summary

Related reading

Tackling System Design Interview Problems

No course covers this one yet

ML System Design practice on Codemia

Passing Arguments with `args`