Call super().build(input_shape) instead of self.built = True in all Keras layers.

Within `build()`, some Keras layers where calling `super().build(input_shape)` while some were calling `self.built = True`. This would result in a different config when serializing whereby layers doing `self.built = True` would not have a `build_config`.

This change makes it consistent between all the layers as well as consistent with Keras 3.

Note that some layers need to call `Layer.build(self, input_shape)` directly to bypass some class' `build()` but still populate the information for the `build_config`.

PiperOrigin-RevId: 678454186

Author: hertschuh

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

@@ -669,10 +669,13 @@ def testStripClusteringSequentialModel(self):
 
     self.assertEqual(self._count_clustered_layers(stripped_model), 0)
     model_config = model.get_config()
+    stripped_model_config = stripped_model.get_config()
+    # New serialization format includes `build_config` in all layers
     for layer in model_config['layers']:
-      # New serialization format includes `build_config` in wrapper
       layer.pop('build_config', None)
-    self.assertEqual(model_config, stripped_model.get_config())
+    for layer in stripped_model_config['layers']:
+      layer.pop('build_config', None)
+    self.assertEqual(model_config, stripped_model_config)
 
   def testClusterStrippingFunctionalModel(self):
     """Verifies that stripping the clustering wrappers from a functional model produces the expected config."""

tensorflow_model_optimization/python/core/common/keras/compression/internal/optimize.py

@@ -404,6 +404,9 @@ def create_layer_for_training(layer, algorithm):
   if compressible_weights:
     # Set pretrained weight values.
     wrapped_layer.build(input_shape)
+    # Clear `_build_input_shape` so that `build()` is not immediately called
+    # during reloading. We want the wrapper layer to initiate `build()`.
+    wrapped_layer.layer._build_input_shape = None  # pylint: disable=protected-access
     training_weights = _map_to_training_weights(
         algorithm,
         layer,
@@ -445,6 +448,9 @@ def create_layer_for_inference(layer: _TrainingWrapper, algorithm):
   layer_for_inference = _InferenceWrapper(cloned_layer, algorithm,
                                           compressible_training_tensors)
   layer_for_inference.build(input_shape)
+  # Clear `_build_input_shape` so that `build()` is not immediately called
+  # during reloading. We want the wrapper layer to initiate `build()`.
+  layer_for_inference.layer._build_input_shape = None  # pylint: disable=protected-access
 
   if layer.get_weights():
     # Set weights of layer for inference according to what was trained.

tensorflow_model_optimization/python/core/sparsity/keras/pruning_wrapper_test.py

@@ -122,17 +122,17 @@ def testPruneModel(self):
 
     # Test serialization
     model_config = self.model.get_config()
-    for layer in model_config['layers']:
-      layer.pop('build_config', None)
-    self.assertEqual(
+    pruned_model = self.model.__class__.from_config(
         model_config,
-        self.model.__class__.from_config(
-            self.model.get_config(),
-            custom_objects={
-                'PruneLowMagnitude': pruning_wrapper.PruneLowMagnitude
-            },
-        ).get_config(),
+        custom_objects={'PruneLowMagnitude': pruning_wrapper.PruneLowMagnitude},
     )
+    pruned_model_config = pruned_model.get_config()
+    # New serialization format includes `build_config` in all layers
+    for layer in model_config['layers']:
+      layer.pop('build_config', None)
+    for layer in pruned_model_config['layers']:
+      layer.pop('build_config', None)
+    self.assertEqual(model_config, pruned_model_config)
 
   def testCustomLayerNonPrunable(self):
     layer = CustomLayer(input_dim=16, output_dim=32)