Fixes improper serialization code route for functional models.

PiperOrigin-RevId: 524097977

official/nlp/modeling/layers/position_embedding.py

@@ -72,7 +72,7 @@ class PositionEmbedding(tf.keras.layers.Layer):
     return dict(list(base_config.items()) + list(config.items()))
 
   def build(self, input_shape):
-    dimension_list = input_shape.as_list()
+    dimension_list = input_shape
     width = dimension_list[-1]
     weight_sequence_length = self._max_length
 

official/nlp/modeling/layers/transformer_scaffold_test.py

@@ -56,7 +56,7 @@ class ValidatedFeedforwardLayer(tf.keras.layers.Layer):
     self.activation = activation
 
   def build(self, input_shape):
-    hidden_size = input_shape.as_list()[-1]
+    hidden_size = input_shape[-1]
     self._feedforward_dense = tf.keras.layers.EinsumDense(
         '...x,xy->...y',
         output_shape=hidden_size,

official/projects/mosaic/qat/modeling/factory.py

@@ -54,6 +54,7 @@ def build_qat_mosaic_model(
       'L2': tf.keras.regularizers.l2,
   }
 
+  model.use_legacy_config = True  # Ensures old Keras serialization format
   # Apply QAT to backbone (a tf.keras.Model) first, and then neck and head.
   with tfmot.quantization.keras.quantize_scope(scope_dict):
     annotated_backbone = tfmot.quantization.keras.quantize_annotate_model(

official/projects/qat/vision/modeling/factory.py

@@ -104,11 +104,13 @@ def build_qat_classification_model(
         return tfmot.quantization.keras.quantize_annotate_layer(layer)
       return layer
 
+    backbone_optimized_model.use_legacy_config = True
     annotated_model = tf.keras.models.clone_model(
         backbone_optimized_model,
         clone_function=apply_quantization_to_dense,
     )
 
+    annotated_model.use_legacy_config = True
     if quantization.change_num_bits:
       optimized_model = tfmot.quantization.keras.quantize_apply(
           annotated_model,
@@ -236,6 +238,7 @@ def build_qat_segmentation_model(
       'L2': tf.keras.regularizers.l2,
   }
 
+  model.use_legacy_config = True  # Ensures old Keras serialization format
   # Apply QAT to backbone (a tf.keras.Model) first.
   with tfmot.quantization.keras.quantize_scope(scope_dict):
     annotated_backbone = tfmot.quantization.keras.quantize_annotate_model(
@@ -259,10 +262,12 @@ def build_qat_segmentation_model(
         return tfmot.quantization.keras.quantize_annotate_layer(layer)
       return layer
 
+    backbone_optimized_model.use_legacy_config = True
     annotated_model = tf.keras.models.clone_model(
         backbone_optimized_model,
         clone_function=apply_quantization_to_layers,
     )
+    annotated_model.use_legacy_config = True
     optimized_model = tfmot.quantization.keras.quantize_apply(
         annotated_model, scheme=schemes.Default8BitQuantizeScheme())
 

official/vision/modeling/layers/nn_blocks_test.py

@@ -431,7 +431,7 @@ class ValidatedFeedforwardLayer(tf.keras.layers.Layer):
     self.activation = activation
 
   def build(self, input_shape):
-    hidden_size = input_shape.as_list()[-1]
+    hidden_size = input_shape[-1]
     self._feedforward_dense = tf.keras.layers.EinsumDense(
         '...x,xy->...y',
         output_shape=hidden_size,