Switches Keras object serialization to new logic and changes public A…

…PI for deserialize_keras_object/serialize_keras_object to the new functions.

PiperOrigin-RevId: 480676373

edward2/tensorflow/constraints.py

@@ -84,7 +84,6 @@ def serialize(initializer):
 def deserialize(config, custom_objects=None):
   return tf.keras.utils.deserialize_keras_object(
       config,
-      module_objects=globals(),
       custom_objects=custom_objects,
       printable_module_name='constraints')
 

edward2/tensorflow/initializers.py

@@ -138,8 +138,7 @@ def __init__(self,
       distribution: Random distribution to use. One of "truncated_normal", or
         "untruncated_normal".
       seed: A Python integer. Used to create random seeds. See
-        `tf.set_random_seed`
-        for behavior.
+        `tf.set_random_seed` for behavior.
 
     Raises:
       ValueError: In case of an invalid value for the "scale", mode" or
@@ -529,7 +528,9 @@ def get_config(self):
 
 
 class TrainableHeNormal(TrainableNormal):
-  """Trainable normal initialized per He et al. 2015, given a ReLU nonlinearity.
+  """Trainable normal initialized per He et al.
+
+  2015, given a ReLU nonlinearity.
 
   The distribution is initialized to a Normal scaled by `sqrt(2 / fan_in)`,
   where `fan_in` is the number of input units. A ReLU nonlinearity is assumed
@@ -857,7 +858,6 @@ def serialize(initializer):
 def deserialize(config, custom_objects=None):
   return tf.keras.utils.deserialize_keras_object(
       config,
-      module_objects=globals(),
       custom_objects=custom_objects,
       printable_module_name='initializers')
 

edward2/tensorflow/regularizers.py

@@ -394,7 +394,6 @@ def serialize(initializer):
 def deserialize(config, custom_objects=None):
   return tf.keras.utils.deserialize_keras_object(
       config,
-      module_objects=globals(),
       custom_objects=custom_objects,
       printable_module_name='regularizers')