Fixes improper serialization code route for functional models.

PiperOrigin-RevId: 524097977

Author: nkovela1

tensorflow_model_optimization/python/core/clustering/keras/cluster_test.py

@@ -686,6 +686,10 @@ def testClusterStrippingFunctionalModel(self):
     stripped_model = cluster.strip_clustering(clustered_model)
 
     self.assertEqual(self._count_clustered_layers(stripped_model), 0)
+
+    # Ensures old Keras serialization format
+    model.use_legacy_config = True
+    stripped_model.use_legacy_config = True
     self.assertEqual(model.get_config(), stripped_model.get_config())
 
   def testClusterWeightsStrippedWeights(self):

tensorflow_model_optimization/python/core/quantization/keras/graph_transformations/model_transformer.py

@@ -554,6 +554,9 @@ def transform(self):
     #   'output_layers': [ ... ],
     #   'name': 'MODEL_NAME',
     #
+
+    # Ensures old Keras serialization format
+    self.model.use_legacy_config = True
     self._config = self.model.get_config()
 
     # Stores map of Transform -> List of layer names matched by transform.
@@ -631,4 +634,6 @@ def transform(self):
         if names_and_weights:
           self._set_layer_names_and_weights(layer, names_and_weights)
 
+    # Ensures old Keras serialization format
+    transformed_model.use_legacy_config = True
     return transformed_model, copy.deepcopy(self._layer_metadata_map)

tensorflow_model_optimization/python/core/quantization/keras/graph_transformations/model_transformer_test.py

@@ -463,6 +463,8 @@ def replacement(self, match_layer):
       inp = keras.layers.Input((3,))
       out = keras.layers.Dense(2, activation='relu')(inp)
       model_fused = keras.Model(inp, out)
+      # Ensures old Keras serialization format
+      model_fused.use_legacy_config = True
     else:
       model_fused = keras.Sequential(
           [keras.layers.Dense(2, activation='relu', input_shape=(3,))])

tensorflow_model_optimization/python/core/quantization/keras/quantize.py

@@ -461,11 +461,15 @@ def _quantize(layer):  # pylint: disable=missing-docstring
         '`quantize_scope` for your calls to `quantize_model` and '
         '`quantize_apply`. [%s].' % er) from er
 
+  if hasattr(model, 'use_legacy_config'):
+    model_copy.use_legacy_config = model.use_legacy_config
+
   # 2. Remove QuantizeAnnotate wrappers from the layers in the model. This
   # extracts the original model structure (easier to transform), and
   # stores relevant quantization information in a map.
-  (unwrapped_model, layer_quantize_map,
-   requires_output_quantize) = _extract_original_model(model_copy)
+  (unwrapped_model, layer_quantize_map, requires_output_quantize) = (
+      _extract_original_model(model_copy)
+  )
   # Model cloning excludes input layers. Add input layers into the map
   # since they need to be matched for patterns as well.
   # pylint: disable=protected-access