Move code to use internal smart_cond

MOT code currently uses `smart_cond` from Keras which is not a
publicly supported function. Due to this, changes in Keras often
lead to code breaking in MOT.

Moving the code to use an internal copy of the code.

PiperOrigin-RevId: 325110684

Author: nutsiepully

tensorflow_model_optimization/python/core/keras/BUILD

@@ -11,6 +11,7 @@ py_library(
     srcs = ["__init__.py"],
     deps = [
         ":compat",
+        ":utils",
     ],
 )
 
@@ -33,3 +34,14 @@ py_library(
         # tensorflow dep1,
     ],
 )
+
+py_library(
+    name = "utils",
+    srcs = ["utils.py"],
+    srcs_version = "PY3",
+    deps = [
+        # python:control_flow_ops tensorflow dep2,
+        # python:smart_cond tensorflow dep2,
+        # python:variables tensorflow dep2,
+    ],
+)

tensorflow_model_optimization/python/core/keras/utils.py

@@ -0,0 +1,54 @@
+# Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Utility functions common to MOT techniques.
+
+`smart_cond` is not exposed as a public TF function with stabilty guarantees,
+and hence changes to it break our code. So, we make a copy of it here.
+"""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+# TODO(b/151772467): Move away from depending on private APIs.
+from tensorflow.python.framework import smart_cond as smart_module
+from tensorflow.python.ops import control_flow_ops
+from tensorflow.python.ops import variables
+
+
+def smart_cond(pred, true_fn=None, false_fn=None, name=None):  # pylint: disable=invalid-name
+  """Return either `true_fn()` if predicate `pred` is true else `false_fn()`.
+
+  If `pred` is a bool or has a constant value, we return either `true_fn()`
+  or `false_fn()`, otherwise we use `tf.cond` to dynamically route to both.
+
+  Arguments:
+    pred: A scalar determining whether to return the result of `true_fn` or
+      `false_fn`.
+    true_fn: The callable to be performed if pred is true.
+    false_fn: The callable to be performed if pred is false.
+    name: Optional name prefix when using `tf.cond`.
+
+  Returns:
+    Tensors returned by the call to either `true_fn` or `false_fn`.
+
+  Raises:
+    TypeError: If `true_fn` or `false_fn` is not callable.
+  """
+  if isinstance(pred, variables.Variable):
+    return control_flow_ops.cond(
+        pred, true_fn=true_fn, false_fn=false_fn, name=name)
+  return smart_module.smart_cond(
+      pred, true_fn=true_fn, false_fn=false_fn, name=name)

tensorflow_model_optimization/python/core/quantization/keras/BUILD

@@ -150,6 +150,7 @@ py_library(
     deps = [
         ":quant_ops",
         # tensorflow dep1,
+        "//tensorflow_model_optimization/python/core/keras:utils",
     ],
 )
 
@@ -178,6 +179,7 @@ py_library(
     deps = [
         ":quantizers",
         # tensorflow dep1,
+        "//tensorflow_model_optimization/python/core/keras:utils",
     ],
 )
 
@@ -208,6 +210,7 @@ py_library(
         ":quantize_config",
         ":quantizers",
         # tensorflow dep1,
+        "//tensorflow_model_optimization/python/core/keras:utils",
     ],
 )
 

tensorflow_model_optimization/python/core/quantization/keras/layers/BUILD

@@ -26,6 +26,7 @@ py_library(
         # python/keras:backend tensorflow dep2,
         # python/keras/layers tensorflow dep2,
         # python/keras/utils:engine_utils tensorflow dep2,
+        "//tensorflow_model_optimization/python/core/keras:utils",
         "//tensorflow_model_optimization/python/core/quantization/keras:quantize_aware_activation",
         "//tensorflow_model_optimization/python/core/quantization/keras:quantizers",
         "//tensorflow_model_optimization/python/core/quantization/keras/default_8bit:default_8bit_quantizers",

tensorflow_model_optimization/python/core/quantization/keras/layers/conv_batchnorm.py

@@ -26,13 +26,14 @@
 from tensorflow.python.keras.layers import convolutional
 from tensorflow.python.keras.layers import deserialize as deserialize_layer
 from tensorflow.python.keras.layers import normalization
-from tensorflow.python.keras.utils import control_flow_util
 from tensorflow.python.keras.utils import conv_utils
 from tensorflow.python.ops import array_ops
 from tensorflow.python.ops import math_ops
 from tensorflow.python.ops import nn
 from tensorflow.python.ops import nn_ops
 
+from tensorflow_model_optimization.python.core.keras import utils
+
 from tensorflow_model_optimization.python.core.quantization.keras import quantizers
 from tensorflow_model_optimization.python.core.quantization.keras.default_8bit import default_8bit_quantizers
 
@@ -81,8 +82,8 @@ def quantizer_fn():
 
       return quantizer_fn
 
-    return control_flow_util.smart_cond(training, make_quantizer_fn(True),
-                                        make_quantizer_fn(False))
+    return utils.smart_cond(
+        training, make_quantizer_fn(True), make_quantizer_fn(False))
 
   def _apply_activation_quantizer(self, training, activation_output):
     """All Keras call() logic for applying weight quantization."""
@@ -101,8 +102,8 @@ def quantizer_fn():
 
       return quantizer_fn
 
-    return control_flow_util.smart_cond(training, make_quantizer_fn(True),
-                                        make_quantizer_fn(False))
+    return utils.smart_cond(
+        training, make_quantizer_fn(True), make_quantizer_fn(False))
 
   @staticmethod
   def _from_config(cls_initializer, config):

tensorflow_model_optimization/python/core/quantization/keras/quantize_aware_activation.py

@@ -20,8 +20,7 @@
 
 import tensorflow as tf
 
-# TODO(b/139939526): move to public API.
-from tensorflow.python.keras.utils import control_flow_util
+from tensorflow_model_optimization.python.core.keras import utils
 
 activations = tf.keras.activations
 
@@ -161,14 +160,14 @@ def quantizer_fn(x=x,
 
     x = inputs
     if self._should_pre_quantize():
-      x = control_flow_util.smart_cond(
+      x = utils.smart_cond(
           self._training, make_quantizer_fn(True, x, self._pre_activation_vars),
           make_quantizer_fn(False, x, self._pre_activation_vars))
 
     x = self.activation(x, *args, **kwargs)
 
     if self._should_post_quantize():
-      x = control_flow_util.smart_cond(
+      x = utils.smart_cond(
           self._training, make_quantizer_fn(True, x,
                                             self._post_activation_vars),
           make_quantizer_fn(False, x, self._post_activation_vars))

tensorflow_model_optimization/python/core/quantization/keras/quantize_layer.py

@@ -23,7 +23,8 @@
 
 import tensorflow as tf
 
-from tensorflow.python.keras.utils import control_flow_util
+from tensorflow_model_optimization.python.core.keras import utils
+
 from tensorflow_model_optimization.python.core.quantization.keras import quantizers
 
 serialize_keras_object = tf.keras.utils.serialize_keras_object
@@ -70,8 +71,8 @@ def quantizer_fn():
 
       return quantizer_fn
 
-    return control_flow_util.smart_cond(training, _make_quantizer_fn(True),
-                                        _make_quantizer_fn(False))
+    return utils.smart_cond(
+        training, _make_quantizer_fn(True), _make_quantizer_fn(False))
 
   def get_config(self):
     base_config = super(QuantizeLayer, self).get_config()

tensorflow_model_optimization/python/core/quantization/keras/quantize_wrapper.py

@@ -28,9 +28,9 @@
 
 import tensorflow as tf
 
-# TODO(b/139939526): move to public API.
-from tensorflow.python.keras.utils import control_flow_util
 from tensorflow.python.util import tf_inspect
+
+from tensorflow_model_optimization.python.core.keras import utils
 from tensorflow_model_optimization.python.core.quantization.keras import quantize_aware_activation
 
 deserialize_keras_object = tf.keras.utils.deserialize_keras_object
@@ -141,7 +141,7 @@ def call(self, inputs, training=None):
 
     quantized_weights = []
     for unquantized_weight, quantizer, quantizer_vars in self._weight_vars:
-      quantized_weight = control_flow_util.smart_cond(
+      quantized_weight = utils.smart_cond(
           training,
           self._make_quantizer_fn(quantizer, unquantized_weight, True,
                                   quantizer_vars),
@@ -175,7 +175,7 @@ def call(self, inputs, training=None):
       raise RuntimeError('Multiple output tensors not handled currently.')
 
     output_quantizer = self._output_quantizers[0]
-    return control_flow_util.smart_cond(
+    return utils.smart_cond(
         training,
         self._make_quantizer_fn(output_quantizer, outputs, True,
                                 self._output_quantizer_vars),

tensorflow_model_optimization/python/core/sparsity/keras/BUILD

@@ -92,8 +92,8 @@ py_library(
         ":pruning_schedule",
         # numpy dep1,
         # tensorflow dep1,
-        # python/keras/utils:control_flow_util tensorflow dep2,
         # python/keras/utils:generic_utils tensorflow dep2,
+        "//tensorflow_model_optimization/python/core/keras:utils",
     ],
 )
 

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

@@ -25,9 +25,10 @@
 import tensorflow as tf
 
 # b/(139939526): update to use public API.
-from tensorflow.python.keras.utils import control_flow_util
 from tensorflow.python.keras.utils import generic_utils
 
+from tensorflow_model_optimization.python.core.keras import utils
+
 from tensorflow_model_optimization.python.core.sparsity.keras import prunable_layer
 from tensorflow_model_optimization.python.core.sparsity.keras import prune_registry
 from tensorflow_model_optimization.python.core.sparsity.keras import pruning_impl
@@ -247,7 +248,7 @@ def add_update():
     def no_op():
       return tf.no_op('no_update')
 
-    update_op = control_flow_util.smart_cond(training, add_update, no_op)
+    update_op = utils.smart_cond(training, add_update, no_op)
     self.add_update(update_op)
     # Always execute the op that performs weights = weights * mask
     # Relies on UpdatePruningStep callback to ensure the weights