Explicitly import estimator from tensorflow as a separate import instead of accessing it via tf.estimator and depend on the tensorflow estimator target.

PiperOrigin-RevId: 437237782

Author: LukeWood

tensorflow_model_optimization/python/core/sparsity/keras/BUILD

@@ -175,6 +175,7 @@ py_strict_library(
     deps = [
         ":pruning_wrapper",
         # tensorflow dep1,
+        # tensorflow:tensorflow_compat_v1_estimator dep1,
         # python/framework:for_generated_wrappers tensorflow dep2,
     ],
 )

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

@@ -15,12 +15,13 @@
 """Utility functions for making pruning wrapper work with estimators."""
 
 import tensorflow.compat.v1 as tf
+from tensorflow.compat.v1 import estimator as tf_estimator
 
 from tensorflow.python.framework import ops
 from tensorflow_model_optimization.python.core.sparsity.keras.pruning_wrapper import PruneLowMagnitude
 
 
-class PruningEstimatorSpec(tf.estimator.EstimatorSpec):
+class PruningEstimatorSpec(tf_estimator.EstimatorSpec):
   """Returns an EstimatorSpec modified to prune the model while training."""
 
   def __new__(cls, model, step=None, train_op=None, **kwargs):

tensorflow_model_optimization/python/examples/sparsity/keras/mnist/BUILD

@@ -23,6 +23,7 @@ py_strict_binary(
         # numpy dep1,
         # six dep1,
         # tensorflow dep1,
+        # tensorflow:tensorflow_compat_v1_estimator dep1,
         "//tensorflow_model_optimization/python/core/sparsity/keras:estimator_utils",
         "//tensorflow_model_optimization/python/core/sparsity/keras:prune",
         "//tensorflow_model_optimization/python/core/sparsity/keras:pruning_schedule",