Remove pruning wildcard imports for determining what symbols are public APIs, to mirror quantization.

This prevents accidentally exposing private symbols as public APIs.

Validated that build_docs.py produces API docs with expected symbols, even without `private_map` blacklist.

PiperOrigin-RevId: 317025882

Author: alanchiao

build_docs.py

@@ -80,13 +80,6 @@ def main(unused_argv):
       code_url_prefix=FLAGS.code_url_prefix,
       search_hints=FLAGS.search_hints,
       site_path=FLAGS.site_path,
-      private_map={
-          "tfmot.sparsity.keras": [
-              # List of internal classes which get exposed when imported.
-              "InputLayer", "custom_object_scope", "pruning_sched",
-              "pruning_wrapper", "absolute_import", "division", "print_function"
-          ]
-      },
   )
 
   doc_generator.build(output_dir=FLAGS.output_dir)

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

@@ -13,9 +13,19 @@
 # limitations under the License.
 # ==============================================================================
 """Module containing sparsity code built on Keras abstractions."""
-# pylint: disable=wildcard-import
+# pylint: disable=g-bad-import-order
+
+from tensorflow_model_optimization.python.core.sparsity.keras.prune import prune_low_magnitude
+from tensorflow_model_optimization.python.core.sparsity.keras.prune import prune_scope
+from tensorflow_model_optimization.python.core.sparsity.keras.prune import strip_pruning
+
+from tensorflow_model_optimization.python.core.sparsity.keras.pruning_callbacks import UpdatePruningStep
+from tensorflow_model_optimization.python.core.sparsity.keras.pruning_callbacks import PruningSummaries
+
+from tensorflow_model_optimization.python.core.sparsity.keras.pruning_schedule import PruningSchedule
+from tensorflow_model_optimization.python.core.sparsity.keras.pruning_schedule import ConstantSparsity
+from tensorflow_model_optimization.python.core.sparsity.keras.pruning_schedule import PolynomialDecay
+
 from tensorflow_model_optimization.python.core.sparsity.keras.prunable_layer import PrunableLayer
-from tensorflow_model_optimization.python.core.sparsity.keras.prune import *
-from tensorflow_model_optimization.python.core.sparsity.keras.pruning_callbacks import *
-from tensorflow_model_optimization.python.core.sparsity.keras.pruning_schedule import *
-# pylint: enable=wildcard-import
+
+# pylint: enable=g-bad-import-order