Introduction

tf.data.Dataset.apply() is often misunderstood as a general-purpose way to reshape dataset elements. In practice, simple reshaping belongs in map, where you transform each element with tf.reshape. apply() is for dataset-level transformations, not for ordinary tensor shape changes inside each record.

Use map To Reshape Elements

If each dataset element is a tensor and you want to reshape that tensor, write a mapping function.

1import tensorflow as tf
2
3
4dataset = tf.data.Dataset.from_tensor_slices(tf.range(12))
5dataset = dataset.batch(4)
6dataset = dataset.map(lambda x: tf.reshape(x, (2, 2)))
7
8for item in dataset:
9    print(item.numpy())

That is the idiomatic solution. Each element produced by batch(4) has shape (4,), and the mapping function reshapes it to (2, 2).

The same pattern works for images, sequences, and label pairs.

1import tensorflow as tf
2
3
4def reshape_example(features, label):
5    features = tf.reshape(features, (28, 28, 1))
6    return features, label
7
8
9images = tf.random.uniform((5, 784))
10labels = tf.constant([0, 1, 2, 3, 4])
11dataset = tf.data.Dataset.from_tensor_slices((images, labels))
12dataset = dataset.map(reshape_example)
13
14for image, label in dataset.take(1):
15    print(image.shape, label.numpy())

That is how you prepare flat vectors for a convolutional model.

What apply() Is Actually For

Dataset.apply() accepts a function that takes a dataset and returns another dataset. That means it operates on the pipeline, not on individual elements.

Conceptually, it looks like this:

new_dataset = dataset.apply(transformation_function)

Historically, TensorFlow used apply() with tf.data.experimental helpers such as bucketing and performance-related transformations. It is not the normal tool for element reshape operations.

A simplified dataset-level transformation looks like this:

1import tensorflow as tf
2
3
4def repeat_twice(ds):
5    return ds.repeat(2)
6
7
8base = tf.data.Dataset.range(3)
9expanded = base.apply(repeat_twice)
10
11for value in expanded:
12    print(int(value))

That demonstrates the right level of abstraction: the function receives the whole dataset pipeline.

Reshape The Tensor, Not The Dataset Container

A common source of confusion is the word "reshape." Usually the real task is reshaping the tensors inside the dataset, not changing the dataset object itself.

For example, if you batch then reshape, the order matters.

1import tensorflow as tf
2
3
4dataset = tf.data.Dataset.range(8)
5dataset = dataset.batch(4)
6dataset = dataset.map(lambda x: tf.reshape(x, (2, 2)))
7
8for item in dataset:
9    print(item.shape)

If you attempted the reshape before batching, you would be reshaping scalar elements instead of vectors of length 4, which is a different problem entirely.

That is why the placement of map in the pipeline matters.

When apply() Still Makes Sense

apply() is still useful when you want to package a reusable dataset transformation pipeline.

1import tensorflow as tf
2
3
4def prepare_batches(batch_size):
5    def transform(ds):
6        return ds.shuffle(100).batch(batch_size).prefetch(tf.data.AUTOTUNE)
7    return transform
8
9
10base = tf.data.Dataset.range(20)
11prepared = base.apply(prepare_batches(5))
12
13for batch in prepared.take(1):
14    print(batch.numpy())

Even here, if you need to reshape each batch, you would still add a map step inside the transformation function.

Common Pitfalls

The biggest mistake is trying to use apply() for ordinary element-wise tensor reshape. For that, use map with tf.reshape.

Another common error is reshaping before batching when the target shape assumes batched data. Check the element shape at each stage of the pipeline.

Developers also sometimes use internal or experimental transformation helpers when a straightforward map, batch, and prefetch chain would be clearer.

Finally, remember that apply() receives a dataset and must return a dataset. A function that expects one tensor element is the wrong shape of function for apply().

Summary

• Use dataset.map(lambda x: tf.reshape(x, ...)) to reshape dataset elements.
• 'Dataset.apply() is for dataset-level transformations, not simple element reshaping.'
• The position of map in the pipeline affects what shape is available to reshape.
• Batch first if the target shape depends on batched dimensions.
• 'apply() is most useful for reusable pipeline transformations.'
• Keep element transforms and dataset transforms conceptually separate.