RFC: Supporting tf.function in tf.transform

[varshaan]

This RFC will be open for comment until Friday, November 6th, 2020.

cc @jsimsa @zoyahav @katsiapis

Status	Proposed
RFC #	308
Author(s)	Varshaa Naganathan ([email protected])
Sponsor	Konstantinos Katsiapis ([email protected]), Jiri Simsa ([email protected]), Zohar Yahav ([email protected])
Updated	2020-10-12

Objective

For tf.transform users who have TF 2 features enabled, tf.transform will
internally use tf.function to trace the user defined preprocessing function.
This is expected to work for the majority of existing use cases with no user
visible impact. However, it is possible that the user defined preprocessing
function is not traceable by tf.function (i.e. not compatible with TF 2 graph
mode execution). To accommodate such cases, we will provide a
force_tf_compat_v1 escape hatch to allow users to fall back to TF 1
compatibility mode. Keeping the legacy behavior will continue to write the
output artifact as a TF 1 SavedModel and is not guaranteed to be compatible with
V2 Trainer.

To support this change, we need to allow tf.transform analyzers to be traceable
by tf.function. Additionally, for tf.transform to support tf.function usage
completely, it needs to handle nested tf.functions that could by themselves
contain tf.transform analyzers. tf.transform currently only traverses a single
graph generated from the user defined preprocessing function. To support nested
tf.functions, tf.transform will need to recursively traverse all the nested
FuncGraphs, identifying tf.transform analyzer nodes and replacing them with
their final computed value.

A side-effect of this proposal is that the output artifact created by
tf.transform when tf.function is used to trace the preprocessing function will
be a V2 SavedModel.

[ematejska]

There has been no feedback, this is now accepted.