Why a model using tf.py_function can not be serialized?

Question

According to the documentation of ty.py_function a model using it can't be serialized.

The body of the function (i.e. func) will not be serialized in a GraphDef. Therefore, you should not use this function if you need to serialize your model and restore it in a different environment.

Why is serialization not possible?

I was looking for an explanation to why this is the case and alternatives to using tf.py_function but did not find helpful ones.

In my specific case I want to use the Keras Tokenizer and its methods expect numpy arrays - so I am calling it using tf.py_function.

Answer 1

In general, TensorFlow graphs contain a description of operations to be performed on tensors. When you do, for example, tf.add(a, b), a new node is added to the graph indicating that an addition operation should be computed on the tensors a and b.

When you use tf.py_function, TensorFlow has to execute some arbitrary Python code that you are giving as a function. Unfortunately, tf.py_function is not able to "translate" arbitrary functions into graph nodes, so there is no way that the operations can be incorporated as such into the graph. One could consider the option of just embedding the Python code itself into the graph, but since TensorFlow works on multiple languages it would not make sense to have nodes with code specialized for one particular language.

In TensorFlow 2.x, where eager execution is enabled by default, there is tf.function to convert particular functions into graphs, for faster execution (see guide), but there is no direct integration of this functionality with tf.py_function as such.