Fix dependencies for __init__.py files related to python module imports and quantization/keras/collaborative_optimizations

PiperOrigin-RevId: 370621369

Author: teijeong

tensorflow_model_optimization/BUILD

@@ -5,24 +5,22 @@
 #  used to optimize machine learning models for deployment and execution.
 #
 #  https://github.com/tensorflow/model-optimization
+load("//tensorflow_model_optimization:tensorflow_model_optimization.bzl", "py_strict_library")
 
 package(default_visibility = ["//visibility:public"])
 
 licenses(["notice"])
 
 exports_files(["LICENSE"])
 
-py_library(
+py_strict_library(
     name = "tensorflow_model_optimization",
     srcs = ["__init__.py"],
     srcs_version = "PY3",
     deps = [
-        "//tensorflow_model_optimization/python",
+        "//tensorflow_model_optimization/python",  # buildcleaner: keep
+        "//tensorflow_model_optimization/python/core:version",
         "//tensorflow_model_optimization/python/core/api",
-        "//tensorflow_model_optimization/python/core/clustering",
-        "//tensorflow_model_optimization/python/core/keras:compat",
-        "//tensorflow_model_optimization/python/core/quantization",
-        "//tensorflow_model_optimization/python/core/sparsity",
     ],
 )
 

tensorflow_model_optimization/__init__.py

@@ -82,9 +82,11 @@ def _ensure_tf_install():  # pylint: disable=g-statement-before-imports
 # To ensure users only access the expected public API, the API structure is
 # created in the `api` directory. Import all api modules.
 from tensorflow_model_optimization.python.core import version
-# pylint: disable=wildcard-import
-from tensorflow_model_optimization.python.core.api import *
-# pylint: enable=wildcard-import
+
+from tensorflow_model_optimization.python.core.api import clustering
+from tensorflow_model_optimization.python.core.api import experimental
+from tensorflow_model_optimization.python.core.api import quantization
+from tensorflow_model_optimization.python.core.api import sparsity
 
 
 # Use sparsity module to fetch the path for the `api` directory.

tensorflow_model_optimization/python/BUILD

@@ -1,12 +1,14 @@
+load("//tensorflow_model_optimization:tensorflow_model_optimization.bzl", "py_strict_library")
+
 package(default_visibility = ["//visibility:public"])
 
 licenses(["notice"])  # Apache 2.0
 
-py_library(
+py_strict_library(
     name = "python",
     srcs = ["__init__.py"],
     srcs_version = "PY3",
     deps = [
-        "//tensorflow_model_optimization/python/core",
+        "//tensorflow_model_optimization/python/core",  # buildcleaner: keep
     ],
 )

tensorflow_model_optimization/python/core/BUILD

@@ -1,22 +1,24 @@
+load("//tensorflow_model_optimization:tensorflow_model_optimization.bzl", "py_strict_library")
+
 package(default_visibility = ["//visibility:public"])
 
 licenses(["notice"])  # Apache 2.0
 
-py_library(
+py_strict_library(
     name = "core",
     srcs = ["__init__.py"],
     srcs_version = "PY3",
     deps = [
-        ":version",
-        "//tensorflow_model_optimization/python/core/clustering",
-        "//tensorflow_model_optimization/python/core/internal",
-        "//tensorflow_model_optimization/python/core/keras",
-        "//tensorflow_model_optimization/python/core/quantization",
-        "//tensorflow_model_optimization/python/core/sparsity",
+        "//tensorflow_model_optimization/python/core/api",  # buildcleaner: keep
+        "//tensorflow_model_optimization/python/core/clustering",  # buildcleaner: keep
+        "//tensorflow_model_optimization/python/core/internal",  # buildcleaner: keep
+        "//tensorflow_model_optimization/python/core/keras",  # buildcleaner: keep
+        "//tensorflow_model_optimization/python/core/quantization",  # buildcleaner: keep
+        "//tensorflow_model_optimization/python/core/sparsity",  # buildcleaner: keep
     ],
 )
 
-py_library(
+py_strict_library(
     name = "version",
     srcs = ["version.py"],
     srcs_version = "PY3",

tensorflow_model_optimization/python/core/api/BUILD

@@ -1,8 +1,10 @@
+load("//tensorflow_model_optimization:tensorflow_model_optimization.bzl", "py_strict_library")
+
 package(default_visibility = ["//visibility:public"])
 
 licenses(["notice"])  # Apache 2.0
 
-py_library(
+py_strict_library(
     name = "api",
     srcs = [
         "__init__.py",
@@ -23,8 +25,31 @@ py_library(
     ],
     srcs_version = "PY3",
     deps = [
-        "//tensorflow_model_optimization/python/core/clustering/keras",
-        "//tensorflow_model_optimization/python/core/quantization/keras",
-        "//tensorflow_model_optimization/python/core/sparsity/keras",
+        "//tensorflow_model_optimization/python/core/clustering/keras:cluster",
+        "//tensorflow_model_optimization/python/core/clustering/keras:cluster_config",
+        "//tensorflow_model_optimization/python/core/clustering/keras:clusterable_layer",
+        "//tensorflow_model_optimization/python/core/clustering/keras:clustering_algorithm",
+        "//tensorflow_model_optimization/python/core/clustering/keras:clustering_callbacks",
+        "//tensorflow_model_optimization/python/core/clustering/keras/experimental",
+        "//tensorflow_model_optimization/python/core/quantization/keras:quantize",
+        "//tensorflow_model_optimization/python/core/quantization/keras:quantize_config",
+        "//tensorflow_model_optimization/python/core/quantization/keras:quantize_layout_transform",
+        "//tensorflow_model_optimization/python/core/quantization/keras:quantize_registry",
+        "//tensorflow_model_optimization/python/core/quantization/keras:quantize_scheme",
+        "//tensorflow_model_optimization/python/core/quantization/keras:quantize_wrapper",
+        "//tensorflow_model_optimization/python/core/quantization/keras:quantizers",
+        "//tensorflow_model_optimization/python/core/quantization/keras/collaborative_optimizations/cluster_preserve:cluster_utils",
+        "//tensorflow_model_optimization/python/core/quantization/keras/collaborative_optimizations/cluster_preserve:default_8bit_cluster_preserve_quantize_scheme",
+        "//tensorflow_model_optimization/python/core/quantization/keras/collaborative_optimizations/prune_preserve:default_8bit_prune_preserve_quantize_scheme",
+        "//tensorflow_model_optimization/python/core/quantization/keras/default_8bit:default_8bit_quantize_layout_transform",
+        "//tensorflow_model_optimization/python/core/quantization/keras/default_8bit:default_8bit_quantize_registry",
+        "//tensorflow_model_optimization/python/core/quantization/keras/default_8bit:default_8bit_quantize_scheme",
+        "//tensorflow_model_optimization/python/core/quantization/keras/default_8bit:default_8bit_transforms",
+        "//tensorflow_model_optimization/python/core/quantization/keras/graph_transformations:model_transformer",
+        "//tensorflow_model_optimization/python/core/quantization/keras/graph_transformations:transforms",
+        "//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_schedule",
     ],
 )

tensorflow_model_optimization/python/core/keras/BUILD

@@ -1,5 +1,9 @@
+# buildifier: disable=same-origin-load
 load("//tensorflow_model_optimization:tensorflow_model_optimization.bzl", "py_strict_library")
 
+# buildifier: disable=same-origin-load
+load("//tensorflow_model_optimization:tensorflow_model_optimization.bzl", "py_strict_test")
+
 package(default_visibility = [
     "//tensorflow_model_optimization:__subpackages__",
 ])
@@ -55,13 +59,14 @@ py_strict_library(
     ],
 )
 
-py_test(
+py_strict_test(
     name = "metrics_test",
     srcs = ["metrics_test.py"],
     python_version = "PY3",
     deps = [
         ":metrics",
         # mock dep1,
         # tensorflow dep1,
+        # python/eager:monitoring tensorflow dep2,
     ],
 )

tensorflow_model_optimization/python/core/keras/testing/BUILD

@@ -1,10 +1,12 @@
+load("//tensorflow_model_optimization:tensorflow_model_optimization.bzl", "py_strict_library")
+
 package(default_visibility = [
     "//tensorflow_model_optimization:__subpackages__",
 ])
 
 licenses(["notice"])  # Apache 2.0
 
-py_library(
+py_strict_library(
     name = "test_utils_mnist",
     testonly = 1,
     srcs = ["test_utils_mnist.py"],

tensorflow_model_optimization/python/core/quantization/keras/collaborative_optimizations/cluster_preserve/BUILD

@@ -21,7 +21,6 @@ py_library(
         "cluster_utils.py",
     ],
     srcs_version = "PY3",
-    visibility = ["//visibility:private"],
     deps = [
         # tensorflow dep1,
         "//tensorflow_model_optimization/python/core/clustering/keras:clustering_registry",
@@ -34,7 +33,6 @@ py_library(
         "cluster_preserve_quantize_registry.py",
     ],
     srcs_version = "PY3",
-    visibility = ["//visibility:private"],
     deps = [
         ":cluster_utils",
         # tensorflow dep1,

tensorflow_model_optimization/python/core/quantization/keras/collaborative_optimizations/prune_preserve/BUILD

@@ -21,7 +21,6 @@ py_library(
         "prune_preserve_quantize_registry.py",
     ],
     srcs_version = "PY3",
-    visibility = ["//visibility:private"],
     deps = [
         # tensorflow dep1,
         "//tensorflow_model_optimization/python/core/quantization/keras/default_8bit:default_8bit_quantize_registry",