Update tf_model_optimization code to use public keras API

PiperOrigin-RevId: 369779844

Author: qlzh727

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

@@ -316,8 +316,7 @@ def from_config(cls, config):
         module_objects=globals(),
         custom_objects=custom_objects)
 
-    from tensorflow.python.keras.layers import deserialize as deserialize_layer  # pylint: disable=g-import-not-at-top
-    layer = deserialize_layer(config.pop('layer'))
+    layer = keras.layers.deserialize(config.pop('layer'))
     config['layer'] = layer
 
     return cls(**config)