Enable the PruningPolicy API

PiperOrigin-RevId: 373693249

Author: tensorflower-gardener

tensorflow_model_optimization/python/core/api/BUILD

@@ -50,6 +50,7 @@ py_strict_library(
         "//tensorflow_model_optimization/python/core/sparsity/keras:prunable_layer",
         "//tensorflow_model_optimization/python/core/sparsity/keras:prune",
         "//tensorflow_model_optimization/python/core/sparsity/keras:pruning_callbacks",
+        "//tensorflow_model_optimization/python/core/sparsity/keras:pruning_policy",
         "//tensorflow_model_optimization/python/core/sparsity/keras:pruning_schedule",
     ],
 )

tensorflow_model_optimization/python/core/api/sparsity/keras/__init__.py

@@ -28,4 +28,7 @@
 
 from tensorflow_model_optimization.python.core.sparsity.keras.prunable_layer import PrunableLayer
 
+from tensorflow_model_optimization.python.core.sparsity.keras.pruning_policy import PruningPolicy
+from tensorflow_model_optimization.python.core.sparsity.keras.pruning_policy import PruneForLatencyOnXNNPack
+
 # pylint: enable=g-bad-import-order

tensorflow_model_optimization/python/core/sparsity/keras/BUILD

@@ -16,6 +16,7 @@ py_strict_library(
         ":prunable_layer",  # buildcleaner: keep
         ":prune",  # buildcleaner: keep
         ":pruning_callbacks",  # buildcleaner: keep
+        ":pruning_policy",  # buildcleaner: keep
         ":pruning_schedule",  # buildcleaner: keep
     ],
 )

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

@@ -47,6 +47,8 @@ class PruningPolicy(abc.ABC):
   ```
 
   You can inherit this class to write your own custom pruning policy.
+
+  The API is experimental and is subject to change.
   """
 
   @abc.abstractmethod