Change the compress, training function input argument format as same as decompress.

PiperOrigin-RevId: 338989166

Author: Xhark

tensorflow_model_optimization/python/core/common/keras/compression/algorithm.py

@@ -62,7 +62,7 @@ def init_training_weights_repr(
       `tf.keras.layers.Layer.add_weight`for each tf.Variable to create.
     """
 
-  def compress(self, training_weights: List[tf.Tensor]) -> List[tf.Tensor]:
+  def compress(self, *training_weights: tf.Tensor) -> List[tf.Tensor]:
     """Define the operations to compress a single weight after training.
 
     'Compress' can refer to making the weight more amenable to compression
@@ -71,22 +71,22 @@ def compress(self, training_weights: List[tf.Tensor]) -> List[tf.Tensor]:
     The default is an identity.
 
     Args:
-      training_weights: tf.Tensors representing all variables used during
+      *training_weights: tf.Tensors representing all variables used during
         training, for a single compressible weight, in the order returned in
         `init_training_weights_repr`.
 
     Returns:
       List of tf.Tensors to set to compressed or more compressible form.
     """
-    return training_weights
+    return list(training_weights)
 
-  def decompress(self, compressed_weights: List[tf.Tensor]) -> tf.Tensor:
+  def decompress(self, *compressed_weights: tf.Tensor) -> tf.Tensor:
     """Define the operations to decompress a single weight’s compressed form during inference.
 
     The default is an identity. TODO(): actually isn't.
 
     Args:
-       compressed_weights: tf.Tensors representing a single weight’s compressed
+       *compressed_weights: tf.Tensors representing a single weight’s compressed
          form, coming from what’s returned in `compress`.
 
     Returns:
@@ -95,14 +95,14 @@ def decompress(self, compressed_weights: List[tf.Tensor]) -> tf.Tensor:
     return compressed_weights[0]
 
   @abc.abstractmethod
-  def training(self, training_weights: List[tf.Tensor]) -> tf.Tensor:
+  def training(self, *training_weights: tf.Tensor) -> tf.Tensor:
     """Define a piece of the forward pass during training, which operates on a single compressible weight.
 
     TODO(tfmot): throw this error.
     The default throws an error when training occurs.
 
     Args:
-       training_weights: tf.Tensors representing any variables used during
+       *training_weights: tf.Tensors representing any variables used during
          training, for a single compressible weight, in the order returned in
          `init_training_weights_repr`.
 

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

@@ -48,10 +48,7 @@ def init_training_weights_repr(
   def decompress(self, u: tf.Tensor, sv: tf.Tensor) -> tf.Tensor:
     return tf.matmul(u, sv)
 
-  def compress(self, training_weights: List[tf.Tensor]) -> List[tf.Tensor]:
-    assert len(training_weights) == 1
-    weight = training_weights[0]
-
+  def compress(self, weight: tf.Tensor) -> List[tf.Tensor]:
     rank = self.params.rank
     s, u, v = tf.linalg.svd(weight)
 
@@ -73,8 +70,8 @@ def compress(self, training_weights: List[tf.Tensor]) -> List[tf.Tensor]:
     return [u, sv]
 
   # TODO(tfmot): remove in this example, which is just post-training.
-  def training(self, training_weights: List[tf.Tensor]) -> tf.Tensor:
-    return training_weights[0]
+  def training(self, weight: tf.Tensor) -> tf.Tensor:
+    return weight
 
 
 # TODO(tfmot): consider if we can simplify `create_model_for_training` and

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

@@ -234,7 +234,7 @@ def testSVD_PreservesPretrainedWeights(self):
 
     # kernel
     algorithm = svd.SVD(params)
-    w1, w2 = algorithm.compress([tf.constant(dense_layer_weights[0])])
+    w1, w2 = algorithm.compress(tf.constant(dense_layer_weights[0]))
     assert (w1 == dense_layer_compressed_weights[0]).numpy().all()
     assert (w2 == dense_layer_compressed_weights[1]).numpy().all()
 

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

@@ -73,9 +73,7 @@ def init_training_weights_repr(
   def decompress(self, u: tf.Tensor, sv: tf.Tensor) -> tf.Tensor:
     return tf.matmul(u, sv)
 
-  def training(self, training_weights: List[tf.Tensor]) -> tf.Tensor:
-    u = training_weights[0]
-    sv = training_weights[1]
+  def training(self, u: tf.Tensor, sv: tf.Tensor) -> tf.Tensor:
     return self.decompress(u, sv)
 
 

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

@@ -119,7 +119,7 @@ def call(self, inputs):
         training_weight_tensors.append(
             _prevent_constant_folding(v.read_value(), inputs))
 
-      weight_tensor = self.algorithm.training(training_weight_tensors)
+      weight_tensor = self.algorithm.training(*training_weight_tensors)
       setattr(self.layer, attr_name, weight_tensor)
 
     # This assumes that all changes to the forward pass happen "prior" to
@@ -186,7 +186,7 @@ def build(self, input_shape):
     self.compressed_weights = {}
     for attr_name in self.training_tensors:
       training_tensors = self.training_tensors[attr_name]
-      compressed_tensors = self.algorithm.compress(training_tensors)
+      compressed_tensors = self.algorithm.compress(*training_tensors)
       weights = []
       for t in compressed_tensors:
         weight = self.add_weight(name='TODO', shape=t.shape)
@@ -297,7 +297,7 @@ def _map_to_inference_weights(training_weights, algorithm, training_tensors):
   layer_weights_i = 0
   for weight in weights:
     if weight in training_tensors:
-      compressed = algorithm.compress(training_tensors[weight])
+      compressed = algorithm.compress(*training_tensors[weight])
       for c in compressed:
         compressed_weights.append(c.numpy())
       layer_weights_i += len(training_tensors[weight])