Move _prevent_constant_folding applied location from after algorithm function to before.

This change makes ReducesTFLiteModelSize tests passed.

PiperOrigin-RevId: 338386984

Author: Xhark

tensorflow_model_optimization/python/core/common/keras/compression/algorithms/BUILD

@@ -14,6 +14,7 @@ py_library(
 
 py_test(
     name = "same_training_and_inference_test",
+    timeout = "long",
     srcs = ["same_training_and_inference_test.py"],
     python_version = "PY3",
     deps = [
@@ -35,6 +36,7 @@ py_library(
 
 py_test(
     name = "different_training_and_inference_test",
+    timeout = "long",
     srcs = ["different_training_and_inference_test.py"],
     python_version = "PY3",
     deps = [

tensorflow_model_optimization/python/core/common/keras/compression/algorithms/different_training_and_inference_test.py

@@ -170,12 +170,7 @@ def testSVD_HasReasonableAccuracy_TF(self):
 
     self.assertGreater(results[1], 0.60)
 
-  # TODO(tfmot): currently fails - didn't hook up constant
-  # folding prevention correctly.
-  def testSVD_ReducesTFLiteModelSize_Fails(self):
-    return
-
-    # pylint: disable=unreachable
+  def testSVD_ReducesTFLiteModelSize(self):
     model = _build_model()
 
     original_saved_model_dir = _save_as_saved_model(model)
@@ -192,7 +187,6 @@ def testSVD_ReducesTFLiteModelSize_Fails(self):
     compressed_size = os.path.getsize(compressed_tflite_file)
 
     self.assertLess(compressed_size, original_size / 6)
-    # pylint: enable=unreachable
 
   def testSVD_HasReasonableAccuracy_TFLite(self):
     model = _build_model()

tensorflow_model_optimization/python/core/common/keras/compression/algorithms/same_training_and_inference_test.py

@@ -170,12 +170,7 @@ def testSVD_HasReasonableAccuracy_TF(self):
 
     self.assertGreater(results[1], 0.60)
 
-  # TODO(tfmot): currently fails - didn't hook up constant
-  # folding prevention correctly.
-  def testSVD_ReducesTFLiteModelSize_Fails(self):
-    return
-
-    # pylint: disable=unreachable
+  def testSVD_ReducesTFLiteModelSize(self):
     model = _build_model()
 
     original_saved_model_dir = _save_as_saved_model(model)
@@ -191,7 +186,6 @@ def testSVD_ReducesTFLiteModelSize_Fails(self):
     compressed_size = os.path.getsize(compressed_tflite_file)
 
     self.assertLess(compressed_size, original_size / 6)
-    # pylint: enable=unreachable
 
   def testSVD_HasReasonableAccuracy_TFLite(self):
     model = _build_model()

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

@@ -112,15 +112,15 @@ def build(self, input_shape):
 
   def call(self, inputs):
     for attr_name in self.compressible_weights:
-      training_weight_tensors = [
-          v.read_value() for v in self.training_weights[attr_name]
-      ]
+      # TODO(tfmot): move constant folding prevention to the inference graph
+      # only, since constant folding won't happen during training.
+      training_weight_tensors = []
+      for v in self.training_weights[attr_name]:
+        training_weight_tensors.append(
+            _prevent_constant_folding(v.read_value(), inputs))
+
       weight_tensor = self.algorithm.training(training_weight_tensors)
-      # TODO(tfmot): move this to the inference graph only, since
-      # constant folding won't happen during training.
-      non_const_foldable_weight_tensor = _prevent_constant_folding(
-          weight_tensor, inputs)
-      setattr(self.layer, attr_name, non_const_foldable_weight_tensor)
+      setattr(self.layer, attr_name, weight_tensor)
 
     # This assumes that all changes to the forward pass happen "prior" to
     # the nested layer's portion of the forward pass. This suffices since
@@ -198,13 +198,12 @@ def call(self, inputs, training=None):
     for attr_name in self.training_tensors:
       # TODO(tfmot): understand how read_value() is converted to
       # inference in TensorFlow Lite.
-      compressed_weight_tensors = [
-          v.read_value() for v in self.compressed_weights[attr_name]
-      ]
+      compressed_weight_tensors = []
+      for v in self.compressed_weights[attr_name]:
+        compressed_weight_tensors.append(
+            _prevent_constant_folding(v.read_value(), inputs))
       weight_tensor = self.algorithm.decompress(*compressed_weight_tensors)
-      non_const_foldable_weight_tensor = _prevent_constant_folding(
-          weight_tensor, inputs)
-      setattr(self.layer, attr_name, non_const_foldable_weight_tensor)
+      setattr(self.layer, attr_name, weight_tensor)
 
     # TODO(tfmot): handle training arg if needed given this is inference only.
     return self.layer.call(inputs)