What is the difference in purpose between tf.py_function and tf.function?

Question

The difference between the two is muddled in my head, notwithstanding the nuances of what is eager and what isn't. From what I gather, the @tf.function decorator has two benefits in that

1. it converts functions into TensorFlow graphs for performance, and
2. allows for a more Pythonic style of coding by interpreting many (but not all) common-place Python operations into tensor operations, e.g. if into tf.cond, etc.

From the definition of tf.py_function, it seems that it does just #2 above. Hence, why bother with tf.py_function when tf.function does the job with a performance improvement to boot and without the inability of the former to serialize?

Answer 1

They do indeed start to resemble each other as they are improved, so it is useful to see where they come from. Initially, the difference was that:

• @tf.function turns python code into a series of TensorFlow graph nodes.
• tf.py_function wraps an existing python function into a single graph node.

This means that tf.function requires your code to be relatively simple while tf.py_function can handle any python code, no matter how complex.

While this line is indeed blurring, with tf.py_function doing more interpretation and tf.function accepting lot's of complex python commands, the general rule stays the same:

• If you have relatively simple logic in your python code, use tf.function.
• When you use complex code, like large external libraries (e.g. connecting to a database, or loading a large external NLP package) use tf.py_function.