[r0.30.0 cherry-pick] Use jax2tf batch polymorphism for the Penguin Flax experimental example. (#3975)

Author: dhruvesh09

tfx/dependencies.py

@@ -189,7 +189,7 @@ def make_extra_packages_examples():
       #     examples/penguin/experimental/penguin_pipeline_sklearn_gcp.py)
       # Required for the experimental tfx/examples using Flax, e.g.,
       # tfx/examples/penguin.
-      'jax>=0.2.12,<0.3',
+      'jax>=0.2.13,<0.3',
       'jaxlib>=0.1.64,<0.2',
       'flax>=0.3.3,<0.4',
       # Required for tfx/examples/penguin/penguin_utils_cloud_tuner.py

tfx/examples/penguin/penguin_utils_base.py

@@ -17,7 +17,7 @@
 Flax.
 """
 
-from typing import List, Optional, Text
+from typing import List, Text
 import tensorflow as tf
 import tensorflow_transform as tft
 
@@ -39,21 +39,19 @@ def transformed_name(key):
 
 
 def make_serving_signatures(model,
-                            tf_transform_features: tft.TFTransformOutput,
-                            serving_batch_size: Optional[int] = None):
+                            tf_transform_features: tft.TFTransformOutput):
   """Returns the serving signatures.
 
   Args:
     model: the model function to apply to the transformed features.
     tf_transform_features: The transformation to apply to the serialized
       tf.Example.
-    serving_batch_size: an optional specification for a concrete serving batch
-      size.
 
   Returns:
     The signatures to use for saving the mode. The 'serving_default' signature
-    will be a concrete function that takes a serialized tf.Example, parses it,
-    transformes the features and then applies the model.
+    will be a concrete function that takes a batch of unspecified length of
+    serialized tf.Example, parses them, transformes the features and
+    then applies the model.
   """
 
   model.tft_layer = tf_transform_features.transform_features_layer()
@@ -72,8 +70,7 @@ def serve_tf_examples_fn(serialized_tf_examples):
   return {
       'serving_default':
           serve_tf_examples_fn.get_concrete_function(
-              tf.TensorSpec(
-                  shape=[serving_batch_size], dtype=tf.string, name='examples'))
+              tf.TensorSpec(shape=[None], dtype=tf.string, name='examples'))
   }
 
 

tfx/examples/penguin/penguin_utils_flax_experimental.py

@@ -131,8 +131,10 @@ def predict(params: _Params, inputs: _InputBatch):
 
   trained_params = optimizer.target
 
-  # Convert the prediction function to TF.
-  tf_fn = jax2tf.convert(predict, with_gradient=False, enable_xla=True)
+  # Convert the prediction function to TF, with a variable batch dimension
+  # for all inputs.
+  tf_fn = jax2tf.convert(predict, with_gradient=False, enable_xla=True,
+                         polymorphic_shapes=(None, '(b, 1)'))
 
   # Create tf.Variables for the parameters. If you want more useful variable
   # names, you can use `tree.map_structure_with_path` from the `dm-tree`
@@ -327,8 +329,6 @@ def run_fn(fn_args: tfx.components.FnArgs):
       steps_per_epoch=fn_args.train_steps,
       eval_steps_per_epoch=fn_args.eval_steps,
       tensorboard_log_dir=fn_args.model_run_dir)
-  # TODO(b/180721874): batch polymorphic model not yet supported.
 
-  signatures = base.make_serving_signatures(model, tf_transform_output,
-                                            serving_batch_size=1)
+  signatures = base.make_serving_signatures(model, tf_transform_output)
   tf.saved_model.save(model, fn_args.serving_model_dir, signatures=signatures)