Creates non-breaking changes where necessary in preparation for switching all of Keras to new serialization format.

PiperOrigin-RevId: 509911382
Change-Id: Idc317015c81814169d4862cfa3149fa8203897c4

Author: nkovela1

tensorflow_compression/python/entropy_models/continuous_batched_test.py

@@ -148,12 +148,12 @@ def test_compression_works_after_serialization(self):
     noisy = uniform_noise.NoisyNormal(loc=.5, scale=8.)
     em = ContinuousBatchedEntropyModel(noisy, 1, compression=True)
     self.assertIsNot(em._quantization_offset, None)
-    json = tf.keras.utils.serialize_keras_object(em)
+    json = tf.keras.utils.legacy.serialize_keras_object(em)
     weights = em.get_weights()
     x = noisy.base.sample([100])
     x_quantized = em.quantize(x)
     x_compressed = em.compress(x)
-    em = tf.keras.utils.deserialize_keras_object(json)
+    em = tf.keras.utils.legacy.deserialize_keras_object(json)
     em.set_weights(weights)
     self.assertAllEqual(em.compress(x), x_compressed)
     self.assertAllEqual(em.decompress(x_compressed, [100]), x_quantized)
@@ -162,12 +162,12 @@ def test_compression_works_after_serialization_no_offset(self):
     noisy = uniform_noise.NoisyNormal(loc=0, scale=5.)
     em = ContinuousBatchedEntropyModel(noisy, 1, compression=True)
     self.assertIs(em._quantization_offset, None)
-    json = tf.keras.utils.serialize_keras_object(em)
+    json = tf.keras.utils.legacy.serialize_keras_object(em)
     weights = em.get_weights()
     x = noisy.base.sample([100])
     x_quantized = em.quantize(x)
     x_compressed = em.compress(x)
-    em = tf.keras.utils.deserialize_keras_object(json)
+    em = tf.keras.utils.legacy.deserialize_keras_object(json)
     em.set_weights(weights)
     self.assertAllEqual(em.compress(x), x_compressed)
     self.assertAllEqual(em.decompress(x_compressed, [100]), x_quantized)

tensorflow_compression/python/layers/gdn.py

@@ -218,7 +218,7 @@ def alpha_parameter(self, value):
     self._check_not_built()
     # This is necessary to make Keras deserialization via __init__ work.
     if isinstance(value, dict):
-      value = tf.keras.utils.deserialize_keras_object(value)
+      value = tf.keras.utils.legacy.deserialize_keras_object(value)
     if value is not None and not callable(value):
       # It's a constant, so keep it in compute_dtype.
       value = tf.convert_to_tensor(value, dtype=self.compute_dtype)
@@ -233,7 +233,7 @@ def beta_parameter(self, value):
     self._check_not_built()
     # This is necessary to make Keras deserialization via __init__ work.
     if isinstance(value, dict):
-      value = tf.keras.utils.deserialize_keras_object(value)
+      value = tf.keras.utils.legacy.deserialize_keras_object(value)
     if value is not None and not callable(value):
       # It's a constant, so keep it in compute_dtype.
       value = tf.convert_to_tensor(value, dtype=self.compute_dtype)
@@ -248,7 +248,7 @@ def gamma_parameter(self, value):
     self._check_not_built()
     # This is necessary to make Keras deserialization via __init__ work.
     if isinstance(value, dict):
-      value = tf.keras.utils.deserialize_keras_object(value)
+      value = tf.keras.utils.legacy.deserialize_keras_object(value)
     if value is not None and not callable(value):
       # It's a constant, so keep it in compute_dtype.
       value = tf.convert_to_tensor(value, dtype=self.compute_dtype)
@@ -263,7 +263,7 @@ def epsilon_parameter(self, value):
     self._check_not_built()
     # This is necessary to make Keras deserialization via __init__ work.
     if isinstance(value, dict):
-      value = tf.keras.utils.deserialize_keras_object(value)
+      value = tf.keras.utils.legacy.deserialize_keras_object(value)
     if value is not None and not callable(value):
       # It's a constant, so keep it in compute_dtype.
       value = tf.convert_to_tensor(value, dtype=self.compute_dtype)
@@ -431,7 +431,7 @@ def try_serialize(parameter, name):
       if parameter is None:
         return None
       try:
-        return tf.keras.utils.serialize_keras_object(parameter)
+        return tf.keras.utils.legacy.serialize_keras_object(parameter)
       except (ValueError, TypeError):  # Should throw TypeError, but doesn't...
         try:
           return float(parameter)
@@ -449,18 +449,22 @@ def try_serialize(parameter, name):
         rectify=self.rectify,
         data_format=self.data_format,
         alpha_parameter=alpha_parameter,
-        beta_parameter=tf.keras.utils.serialize_keras_object(
-            self.beta_parameter),
-        gamma_parameter=tf.keras.utils.serialize_keras_object(
-            self.gamma_parameter),
+        beta_parameter=tf.keras.utils.legacy.serialize_keras_object(
+            self.beta_parameter
+        ),
+        gamma_parameter=tf.keras.utils.legacy.serialize_keras_object(
+            self.gamma_parameter
+        ),
         epsilon_parameter=epsilon_parameter,
         alpha_initializer=tf.keras.initializers.serialize(
-            self.alpha_initializer),
-        beta_initializer=tf.keras.initializers.serialize(
-            self.beta_initializer),
+            self.alpha_initializer
+        ),
+        beta_initializer=tf.keras.initializers.serialize(self.beta_initializer),
         gamma_initializer=tf.keras.initializers.serialize(
-            self.gamma_initializer),
+            self.gamma_initializer
+        ),
         epsilon_initializer=tf.keras.initializers.serialize(
-            self.epsilon_initializer),
+            self.epsilon_initializer
+        ),
     )
     return config

tensorflow_compression/python/layers/parameters_test.py

@@ -31,9 +31,9 @@ def test_name_and_value_are_reproduced_after_serialization(self):
     parameter = self.cls(initial_value, **self.kwargs)
     name_before = parameter.name
     value_before = parameter()
-    json = tf.keras.utils.serialize_keras_object(parameter)
+    json = tf.keras.utils.legacy.serialize_keras_object(parameter)
     weights = parameter.get_weights()
-    parameter = tf.keras.utils.deserialize_keras_object(json)
+    parameter = tf.keras.utils.legacy.deserialize_keras_object(json)
     self.assertIsInstance(parameter, self.cls)
     self.assertEqual(name_before, parameter.name)
     parameter.set_weights(weights)

tensorflow_compression/python/layers/signal_conv.py

@@ -480,7 +480,7 @@ def kernel_parameter(self, value):
     self._check_not_built()
     # This is necessary to make Keras deserialization via __init__ work.
     if isinstance(value, dict):
-      value = tf.keras.utils.deserialize_keras_object(value)
+      value = tf.keras.utils.legacy.deserialize_keras_object(value)
     if isinstance(value, str):
       if value not in ("variable", "rdft"):
         raise ValueError(f"Unsupported value for kernel_parameter: '{value}'.")
@@ -498,7 +498,7 @@ def bias_parameter(self, value):
     self._check_not_built()
     # This is necessary to make Keras deserialization via __init__ work.
     if isinstance(value, dict):
-      value = tf.keras.utils.deserialize_keras_object(value)
+      value = tf.keras.utils.legacy.deserialize_keras_object(value)
     if isinstance(value, str):
       if value != "variable":
         raise ValueError(f"Unsupported value for bias_parameter: '{value}'.")
@@ -991,7 +991,7 @@ def try_serialize(parameter, name):
       if isinstance(parameter, str):
         return parameter
       try:
-        return tf.keras.utils.serialize_keras_object(parameter)
+        return tf.keras.utils.legacy.serialize_keras_object(parameter)
       except (ValueError, TypeError):  # Should throw TypeError, but doesn't...
         if isinstance(parameter, tf.Variable):
           return "variable"