tf.py_function is only for Eager Mode?

Question

Is tf.py_function only for Eager mode, and not for Graph mode?

According to tf.py_function says, it gives the impression that tf.py_function is to be used at Eager mode.

Wraps a python function into a TensorFlow op that executes it eagerly.

This function allows expressing computations in a TensorFlow graph as Python functions. In particular, it wraps a Python function func in a once-differentiable TensorFlow operation that executes it with eager execution enabled.

You can also use tf.py_function to debug your models at runtime using Python tools, i.e., you can isolate portions of your code that you want to debug, wrap them in Python functions and insert pdb tracepoints or print statements as desired, and wrap those functions in tf.py_function.

Calling tf.py_function will acquire the Python Global Interpreter Lock (GIL) that allows only one thread to run at any point in time. This will preclude efficient parallelization and distribution of the execution of the program.

I believe pdb works with Python interpreter and GIL is inside it, but there is no Python interpreter executing Python code while executing TensorFlow Graph in Graph mode. Hence I believe it is only for Eager Mode. Please confirm if this is correct.

Answer 1

tf.py_function is traced into a tf.Graph node and tf.Graph execution runs it eagerly using Python interpreter. Hence, while TenforFlow is executing in Graph mode, it can also execute a node eagerly.