From 2b8043872cd432866eabaebe7317b0490ca7a16a Mon Sep 17 00:00:00 2001
From: Neel Kovelamudi <nkovela@google.com>
Date: Thu, 13 Apr 2023 14:07:43 -0700
Subject: [PATCH] Fixes improper serialization code route for functional
 models.

PiperOrigin-RevId: 524097977
---
 official/nlp/modeling/layers/position_embedding.py        | 2 +-
 official/nlp/modeling/layers/transformer_scaffold_test.py | 2 +-
 official/projects/mosaic/qat/modeling/factory.py          | 1 +
 official/projects/qat/vision/modeling/factory.py          | 5 +++++
 official/vision/modeling/layers/nn_blocks_test.py         | 2 +-
 5 files changed, 9 insertions(+), 3 deletions(-)

diff --git a/official/nlp/modeling/layers/position_embedding.py b/official/nlp/modeling/layers/position_embedding.py
index dc3ecfcce..0b13fa4cc 100644
--- a/official/nlp/modeling/layers/position_embedding.py
+++ b/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
 
diff --git a/official/nlp/modeling/layers/transformer_scaffold_test.py b/official/nlp/modeling/layers/transformer_scaffold_test.py
index 432dfaa1d..95627dfc6 100644
--- a/official/nlp/modeling/layers/transformer_scaffold_test.py
+++ b/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,
diff --git a/official/projects/mosaic/qat/modeling/factory.py b/official/projects/mosaic/qat/modeling/factory.py
index 1ac0b3a36..57437c510 100644
--- a/official/projects/mosaic/qat/modeling/factory.py
+++ b/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(
diff --git a/official/projects/qat/vision/modeling/factory.py b/official/projects/qat/vision/modeling/factory.py
index 0241b0099..f4d02c7df 100644
--- a/official/projects/qat/vision/modeling/factory.py
+++ b/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())
 
diff --git a/official/vision/modeling/layers/nn_blocks_test.py b/official/vision/modeling/layers/nn_blocks_test.py
index ee2843174..e9f6f39a1 100644
--- a/official/vision/modeling/layers/nn_blocks_test.py
+++ b/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,
-- 
GitLab