Introduction

tf.data.Dataset.map still works in eager mode, but eager execution does not mean the map function becomes an ordinary unrestricted Python callback. The function is still part of the TensorFlow input pipeline, so it works best when written with TensorFlow operations on tensors rather than arbitrary Python-side logic. Most confusion around Dataset.map in eager mode comes from expecting a normal Python loop model instead of a TensorFlow data pipeline model.

Basic map Usage in Eager Mode

A simple example works exactly as you would expect:

1import tensorflow as tf
2
3dataset = tf.data.Dataset.range(5)
4dataset = dataset.map(lambda x: x * 2)
5
6for item in dataset:
7    print(item.numpy())

Because eager execution is enabled by default in modern TensorFlow, iterating over the dataset prints concrete values immediately.

map_func Still Receives Tensors

Even in eager mode, the function passed to map receives tensors, not plain Python integers or strings.

1import tensorflow as tf
2
3def transform(x):
4    print(type(x))
5    return x + 1
6
7for item in tf.data.Dataset.range(3).map(transform):
8    print(item.numpy())

The input is a TensorFlow tensor object. That means TensorFlow ops such as tf.cast, tf.reshape, and tf.strings are usually the right tools inside map_func.

Prefer TensorFlow Ops Inside map

A robust map function uses TensorFlow-native operations so the pipeline stays compatible with graph execution, optimization, and parallelism.

1import tensorflow as tf
2
3def preprocess(x):
4    x = tf.cast(x, tf.float32)
5    return x / 10.0
6
7for item in tf.data.Dataset.range(3).map(preprocess):
8    print(item.numpy())

This is the recommended style even when you are currently running eagerly.

Use tf.py_function Only for Python-Only Logic

If the transformation truly depends on Python code or a non-TensorFlow library, wrap that part with tf.py_function.

1import tensorflow as tf
2
3
4def py_double(x):
5    return x * 2
6
7
8def wrapped(x):
9    y = tf.py_function(py_double, [x], Tout=tf.int64)
10    y.set_shape(x.shape)
11    return y
12
13for item in tf.data.Dataset.range(3).map(wrapped):
14    print(item.numpy())

This works, but it comes with tradeoffs:

• less optimization opportunity
• weaker portability
• shape information often needs manual restoration

So tf.py_function is a fallback, not the default design.

Debug with tf.print, Not Only print

Because Dataset.map may be traced or optimized internally, tf.print is often more reliable than ordinary print for seeing values during mapping.

1import tensorflow as tf
2
3
4def inspect(x):
5    tf.print("value in pipeline:", x)
6    return x
7
8for item in tf.data.Dataset.range(3).map(inspect):
9    pass

This is especially useful when code later runs in less purely eager settings.

Keep the Pipeline Functional

A good map_func should be stateless and deterministic unless there is a strong reason otherwise. The cleaner the function, the easier it is for TensorFlow to parallelize and optimize the pipeline.

That usually means:

• transform input tensors into output tensors
• avoid hidden mutable Python state
• keep side effects minimal

When map starts doing a lot of Python-side work, the pipeline usually becomes slower and harder to reason about.

Common Pitfalls

• Assuming eager mode turns Dataset.map into a normal Python callback with no TensorFlow constraints.
• Writing map_func with ordinary Python code when TensorFlow ops would be more appropriate.
• Using print for debugging and then being confused when tracing or pipeline behavior hides what is happening.
• Reaching for tf.py_function too early instead of keeping the transform in TensorFlow ops.
• Forgetting to restore shape information after tf.py_function.

Summary

• 'Dataset.map works in eager mode, but the mapped function still operates in a TensorFlow data-pipeline context.'
• Write map_func with TensorFlow tensor operations whenever possible.
• Use tf.py_function only for transformations that truly require Python-side execution.
• Prefer tf.print for debugging values inside the pipeline.
• Eager mode makes iteration easier, but it does not remove the structural rules of tf.data pipelines.