PyCharm skips breakpoints or does not step into functions when using TensorFlow map function

Problem

When using PyCharm to debug code with TensorFlow's tf.data.Dataset.map function, breakpoints set in the transform function are skipped. Instead, the debugger might halt in an autograph-generated file in the /tmp/__autograph_generated_fileXXXXX.py directory. This issue occurs even when enabling tf.config.run_functions_eagerly(True) or tf.data.experimental.enable_debug_mode().

Cause

TensorFlow's AutoGraph rewrites Python functions (e.g., those passed to .map) into generated files, and the original source file's breakpoints do not match those in the rewritten version. As a result, the debugger fails to stop at intended locations within the original code. Instead, breakpoint() halts within the temporary generated file.

Resolution

To address this behavior, apply one of these solutions:

1. Debug Using enable_debug_mode() with Debugger Registration

Enable debug mode for TensorFlow and attach the debugger to the worker thread used for running the callback:

import tensorflow as tf

tf.data.experimental.enable_debug_mode()  # Enable debug mode before dataset creation

# Define your transformation function
def transform(x):
    import debugpy  # Handles debugger registration
    debugpy.debug_this_thread()
    y = 20
    return x + y

dataset = tf.data.Dataset.range(5).map(transform)

# Iterate and print results
for val in dataset:
    print(val.numpy())

This approach ensures the debugger recognizes and suspends the internal worker thread, allowing you to debug the transformation code more effectively.

2. Debug Without Using .map Wrapping

As the .map wrapper introduces tracing, you can test functions separately by calling them directly (e.g., transform(tf.constant(0))) to debug their behavior without involving TensorFlow's tracing mechanisms.

3. Monkey-Patch TensorFlow's map

For debugging multiple functions with minimal code changes, consider monkey-patching the tf.data.Dataset.map method within your debugging setup block:

if eval(os.environ.get('PYCHARM_LOCAL_DEBUG', 'False')):
    import pydevd_pycharm
    pydevd_pycharm.settrace('localhost', port=2345, stdout_to_server=True, stderr_to_server=True)

    import tensorflow as tf
    tf.data.experimental.enable_debug_mode()   # Enable debug mode before dataset creation

    import functools, threading, pydevd
    _attached = threading.local()
    _orig_map = tf.data.Dataset.map

    def _debug_map(self, map_func, *args, **kwargs):
        @functools.wraps(map_func)
        def wrapped(*a, **kw):
            if not getattr(_attached, 'done', False):
                pydevd.settrace(suspend=False)
                _attached.done = True
            return map_func(*a, **kw)
        return _orig_map(self, wrapped, *args, **kwargs)

    tf.data.Dataset.map = _debug_map

This generic patch registers any worker thread running the .map callback with the active PyCharm debugger session. It also supports debugging nested .map calls.

Each method should first be tested in isolation on a small pipeline to verify compatibility with your TensorFlow version before full application in extensive pipelines.

Additional Notes

• Ensure compatibility with your TensorFlow version before applying these workarounds.