Create a local copy of keras.utils.generic_utils.to_snake_case at pruning_wrapper.py.

`to_snake_case` is private to the `keras.utils` module and is not exported externally.
Since this is a short utility function, we can simply have a copy of this function locally.

PiperOrigin-RevId: 524779757

Author: dansuh17

tensorflow_model_optimization/python/core/sparsity/keras/BUILD

@@ -105,7 +105,6 @@ py_strict_library(
         ":pruning_impl",
         ":pruning_schedule",
         ":pruning_utils",
-        # keras/utils:generic_utils dep1,
         # numpy dep1,
         # tensorflow:tensorflow_no_contrib dep1,
         "//tensorflow_model_optimization/python/core/keras:compat",

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

@@ -20,10 +20,10 @@
 from __future__ import print_function
 
 import inspect
+import re
 
 # import g3
 
-from keras.utils import generic_utils
 import numpy as np
 import tensorflow as tf
 
@@ -42,6 +42,26 @@
 Wrapper = keras.layers.Wrapper
 
 
+def _to_snake_case(name: str) -> str:
+  """Converts `name` to snake case.
+
+  Example: "TensorFlow" -> "tensor_flow"
+
+  Args:
+    name: The name of some python class.
+
+  Returns:
+    `name` converted to snake case.
+  """
+  intermediate = re.sub('(.)([A-Z][a-z]+)', r'\1_\2', name)
+  insecure = re.sub('([a-z])([A-Z])', r'\1_\2', intermediate).lower()
+  # If the class is private the name starts with "_" which is not secure
+  # for creating scopes. We prefix the name with "private" in this case.
+  if insecure[0] != '_':
+    return insecure
+  return 'private' + insecure
+
+
 class PruneLowMagnitude(Wrapper):
   """This wrapper augments a keras layer so the weight tensor may be pruned.
 
@@ -154,8 +174,9 @@ def __init__(self,
     # TODO(pulkitb): This should be pushed up to the wrappers.py
     # Name the layer using the wrapper and underlying layer name.
     # Prune(Dense) becomes prune_dense_1
-    kwargs.update({'name': '{}_{}'.format(
-        generic_utils.to_snake_case(self.__class__.__name__), layer.name)})
+    kwargs.update(
+        {'name': f'{_to_snake_case(self.__class__.__name__)}_{layer.name}'}
+    )
 
     if isinstance(layer, prunable_layer.PrunableLayer) or hasattr(
         layer, 'get_prunable_weights'):

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

@@ -144,7 +144,14 @@ def testCustomLayerPrunable(self):
     layer = CustomLayerPrunable(input_dim=16, output_dim=32)
     inputs = keras.layers.Input(shape=(16))
     _ = layer(inputs)
-    pruning_wrapper.PruneLowMagnitude(layer, block_pooling_type='MAX')
+    pruned_layer = pruning_wrapper.PruneLowMagnitude(
+        layer, block_pooling_type='MAX'
+    )
+    # The name is the layer's name prefixed by the snake_case version of the
+    # `PruneLowMagnitude` class's name.
+    self.assertEqual(
+        pruned_layer.name, 'prune_low_magnitude_custom_layer_prunable'
+    )
 
   def testCollectPrunableLayers(self):
     lstm_layer = keras.layers.RNN(