Changes to tensor_encoding to make code compatible with both TF 1.x and TF 2.0.

PiperOrigin-RevId: 265079527

Author: PraChetit

tensorflow_model_optimization/python/core/internal/tensor_encoding/core/core_encoder.py

@@ -150,14 +150,14 @@ def initial_state(self, name=None):
       keys as `self.children`, each of which maps to an object like this one,
       recursively.
     """
-    with tf.name_scope(name, 'encoder_initial_state'):
+    with tf.compat.v1.name_scope(name, 'encoder_initial_state'):
       return self._initial_state_impl()
 
   def _initial_state_impl(self):
     """Implementation for the `initial_state` method."""
     children_state = {}
     for key, encoder in six.iteritems(self.children):
-      with tf.name_scope(None, '/'.join([self.stage.name, key])):
+      with tf.compat.v1.name_scope(None, '/'.join([self.stage.name, key])):
         children_state[key] = encoder._initial_state_impl()  # pylint: disable=protected-access
     return {
         EncoderKeys.STATE: self.stage.initial_state(),
@@ -185,14 +185,14 @@ def update_state(self, state, state_update_tensors, name=None):
       recursively.
     """
     values = tf.nest.flatten(state) + tf.nest.flatten(state_update_tensors)
-    with tf.name_scope(name, 'encoder_update_state', values):
+    with tf.compat.v1.name_scope(name, 'encoder_update_state', values):
       return self._update_state_impl(state, state_update_tensors)
 
   def _update_state_impl(self, state, state_update_tensors):
     """Implementation for the `update_state` method."""
     children_states = {}
     for key, encoder in six.iteritems(self.children):
-      with tf.name_scope(None, '/'.join([self.stage.name, key])):
+      with tf.compat.v1.name_scope(None, '/'.join([self.stage.name, key])):
         children_states[key] = encoder._update_state_impl(  # pylint: disable=protected-access
             state[EncoderKeys.CHILDREN][key],
             state_update_tensors[EncoderKeys.CHILDREN][key])
@@ -222,7 +222,8 @@ def get_params(self, state, name=None):
       `self.children`, each of which maps to an object like this one,
       recursively.
     """
-    with tf.name_scope(name, 'encoder_get_params', tf.nest.flatten(state)):
+    with tf.compat.v1.name_scope(name, 'encoder_get_params',
+                                 tf.nest.flatten(state)):
       return self._get_params_impl(state)
 
   def _get_params_impl(self, state):
@@ -234,7 +235,7 @@ def _get_params_impl(self, state):
     children_encode_params = {}
     children_decode_params = {}
     for key, encoder in six.iteritems(self.children):
-      with tf.name_scope(None, '/'.join([self.stage.name, key])):
+      with tf.compat.v1.name_scope(None, '/'.join([self.stage.name, key])):
         children_encode_params[key], children_decode_params[key] = (
             encoder._get_params_impl(state[EncoderKeys.CHILDREN][key]))  # pylint: disable=protected-access
     encode_params[EncoderKeys.CHILDREN] = children_encode_params
@@ -272,8 +273,8 @@ def encode(self, x, encode_params, name=None):
         dictionary can be either `Tensor` objects, non-TensorFlow constants such
         as a `list` or numpy value, or `None`, if the shape is not needed.
     """
-    with tf.name_scope(name, 'encoder_encode',
-                       [x] + tf.nest.flatten(encode_params)):
+    with tf.compat.v1.name_scope(name, 'encoder_encode',
+                                 [x] + tf.nest.flatten(encode_params)):
       return self._encode_impl(x, encode_params)
 
   def _encode_impl(self, x, encode_params):
@@ -290,7 +291,7 @@ def _encode_impl(self, x, encode_params):
     children_state_update_tensors = {}
     children_shapes = {}
     for key, encoder in six.iteritems(self.children):
-      with tf.name_scope(None, '/'.join([self.stage.name, key])):
+      with tf.compat.v1.name_scope(None, '/'.join([self.stage.name, key])):
         (encoded_tensors[key], children_state_update_tensors[key],
          children_shapes[key]) = encoder._encode_impl(  # pylint: disable=protected-access
              encoded_tensors[key], encode_params[EncoderKeys.CHILDREN][key])
@@ -319,7 +320,7 @@ def decode(self, encoded_tensors, decode_params, shape, name=None):
     values = (
         tf.nest.flatten(shape) + tf.nest.flatten(decode_params) +
         tf.nest.flatten(encoded_tensors))
-    with tf.name_scope(name, 'encoder_decode', values):
+    with tf.compat.v1.name_scope(name, 'encoder_decode', values):
       # Calling _decode_before_sum_impl with force_decode=True will decode the
       # entire tree, regardless of potential commutativity with sum.
       return self._decode_before_sum_impl(
@@ -354,7 +355,7 @@ def decode_before_sum(self, encoded_tensors, decode_params, shape, name=None):
     values = (
         tf.nest.flatten(shape) + tf.nest.flatten(decode_params) +
         tf.nest.flatten(encoded_tensors))
-    with tf.name_scope(name, 'encoder_decode_before_sum', values):
+    with tf.compat.v1.name_scope(name, 'encoder_decode_before_sum', values):
       return self._decode_before_sum_impl(
           encoded_tensors, decode_params, shape, force_decode=False)
 
@@ -383,7 +384,7 @@ def _decode_before_sum_impl(self, encoded_tensors, decode_params, shape,
     force_decode |= not self.stage.commutes_with_sum
     for key, value in six.iteritems(encoded_tensors):
       if key in self.children:
-        with tf.name_scope(None, '/'.join([self.stage.name, key])):
+        with tf.compat.v1.name_scope(None, '/'.join([self.stage.name, key])):
           temp_encoded_tensors[key] = (
               self.children[key]._decode_before_sum_impl(  # pylint: disable=protected-access
                   value,
@@ -435,7 +436,7 @@ def decode_after_sum(self,
         tf.nest.flatten(shape) + tf.nest.flatten(decode_params) +
         tf.nest.flatten(encoded_tensors) + [num_summands])
     num_summands = tf.convert_to_tensor(num_summands)
-    with tf.name_scope(name, 'encoder_decode_after_sum', values):
+    with tf.compat.v1.name_scope(name, 'encoder_decode_after_sum', values):
       return self._decode_after_sum_impl(encoded_tensors, decode_params,
                                          num_summands, shape)
 
@@ -450,7 +451,7 @@ def _decode_after_sum_impl(self, encoded_tensors, decode_params, num_summands,
     temp_encoded_tensors = {}
     for key, value in six.iteritems(encoded_tensors):
       if key in self.children:
-        with tf.name_scope(None, '/'.join([self.stage.name, key])):
+        with tf.compat.v1.name_scope(None, '/'.join([self.stage.name, key])):
           temp_encoded_tensors[key] = self.children[key]._decode_after_sum_impl(  # pylint: disable=protected-access
               value, decode_params[EncoderKeys.CHILDREN][key], num_summands,
               shape[EncoderKeys.CHILDREN][key])

tensorflow_model_optimization/python/core/internal/tensor_encoding/core/core_encoder_test.py

@@ -51,6 +51,10 @@
 AN_NORM_UPDATE = test_utils.AdaptiveNormalizeEncodingStage.NORM_STATE_UPDATE_KEY
 
 
+if tf.executing_eagerly():
+  tf.compat.v1.disable_eager_execution()
+
+
 class EncoderTest(tf.test.TestCase):
 
   def test_correct_structure(self):

tensorflow_model_optimization/python/core/internal/tensor_encoding/core/encoding_stage.py

@@ -631,7 +631,7 @@ def _tf_style_initial_state(initial_state_fn):
 
   def actual_initial_state_fn(self, name=None):
     """Modified `initial_state` method."""
-    with tf.name_scope(name, self.name + INITIAL_STATE_SCOPE_SUFFIX):
+    with tf.compat.v1.name_scope(name, self.name + INITIAL_STATE_SCOPE_SUFFIX):
       return initial_state_fn(self, name=name)
 
   return actual_initial_state_fn
@@ -643,7 +643,8 @@ def _tf_style_update_state(update_state_fn):
   def actual_initial_state_fn(self, state, state_update_tensors, name=None):
     """Modified `update_state` method."""
     values = list(state.values()) + list(state_update_tensors.values())
-    with tf.name_scope(name, self.name + UPDATE_STATE_SCOPE_SUFFIX, values):
+    with tf.compat.v1.name_scope(name, self.name + UPDATE_STATE_SCOPE_SUFFIX,
+                                 values):
       state = tf.nest.map_structure(tf.convert_to_tensor, state)
       state_update_tensors = tf.nest.map_structure(tf.convert_to_tensor,
                                                    state_update_tensors)
@@ -657,7 +658,7 @@ def _tf_style_get_params(get_params_fn):
 
   def actual_get_params_fn(self, name=None):
     """Modified `get_params` method."""
-    with tf.name_scope(name, self.name + GET_PARAMS_SCOPE_SUFFIX):
+    with tf.compat.v1.name_scope(name, self.name + GET_PARAMS_SCOPE_SUFFIX):
       return get_params_fn(self, name=name)
 
   return actual_get_params_fn
@@ -668,8 +669,8 @@ def _tf_style_adaptive_get_params(get_params_fn):
 
   def actual_get_params_fn(self, state, name=None):
     """Modified `get_params` method."""
-    with tf.name_scope(name, self.name + GET_PARAMS_SCOPE_SUFFIX,
-                       state.values()):
+    with tf.compat.v1.name_scope(name, self.name + GET_PARAMS_SCOPE_SUFFIX,
+                                 state.values()):
       state = tf.nest.map_structure(tf.convert_to_tensor, state)
       return get_params_fn(self, state, name=name)
 
@@ -682,7 +683,8 @@ def _tf_style_encode(encode_fn):
   def actual_encode_fn(self, x, encode_params, name=None):
     """Modified `encode` method."""
     values = list(encode_params.values()) + [x]
-    with tf.variable_scope(name, self.name + ENCODE_SCOPE_SUFFIX, values):
+    with tf.compat.v1.variable_scope(name, self.name + ENCODE_SCOPE_SUFFIX,
+                                     values):
       x = tf.convert_to_tensor(x)
       encode_params = tf.nest.map_structure(tf.convert_to_tensor, encode_params)
       return encode_fn(self, x, encode_params, name=name)
@@ -701,9 +703,10 @@ def actual_decode_fn(self,
                        name=None):
     """Modified `decode` method."""
     values = list(encoded_tensors.values()) + list(decode_params.values())
-    with tf.variable_scope(name, self.name + DECODE_SCOPE_SUFFIX, values):
+    with tf.compat.v1.variable_scope(name, self.name + DECODE_SCOPE_SUFFIX,
+                                     values):
       encoded_tensors = tf.nest.map_structure(tf.convert_to_tensor,
-                                           encoded_tensors)
+                                              encoded_tensors)
       decode_params = tf.nest.map_structure(tf.convert_to_tensor, decode_params)
       if shape is not None:
         shape = tf.convert_to_tensor(shape)

tensorflow_model_optimization/python/core/internal/tensor_encoding/core/encoding_stage_test.py

@@ -25,6 +25,10 @@
 from tensorflow_model_optimization.python.core.internal.tensor_encoding.testing import test_utils
 
 
+if tf.executing_eagerly():
+  tf.compat.v1.disable_eager_execution()
+
+
 class TFStyleEncodeDecodeTest(tf.test.TestCase, parameterized.TestCase):
   """Tests for `_tf_style_*` decorators.
 
@@ -95,7 +99,7 @@ def test_initial_state_decorator(self, name):
     initial_state = self.evaluate(test_initial_state_fn(stage, name))
 
     # The graph should contain a single node.
-    graph = tf.get_default_graph()
+    graph = tf.compat.v1.get_default_graph()
     self.assertLen(graph.as_graph_def().node, 1)
     if name is not None:
       self._assert_all_graph_nodes_in_name_scope(graph, name)
@@ -117,7 +121,7 @@ def test_update_state_decorator(self, name):
 
     # The graph should contain three nodes. Two for the constants created, and
     # one for their addition.
-    graph = tf.get_default_graph()
+    graph = tf.compat.v1.get_default_graph()
     self.assertLen(graph.as_graph_def().node, 3)
     if name is not None:
       self._assert_all_graph_nodes_in_name_scope(graph, name)
@@ -136,7 +140,7 @@ def test_get_params_decorator(self, name):
         test_get_params_fn(stage, name))
 
     # The graph should contain a single node.
-    graph = tf.get_default_graph()
+    graph = tf.compat.v1.get_default_graph()
     self.assertLen(graph.as_graph_def().node, 1)
     if name is not None:
       self._assert_all_graph_nodes_in_name_scope(graph, name)
@@ -158,7 +162,7 @@ def test_adaptive_get_params_decorator(self, name):
 
     # The graph should contain three nodes. Two for the constants created, and
     # one for the multiplication to create the params.
-    graph = tf.get_default_graph()
+    graph = tf.compat.v1.get_default_graph()
     self.assertLen(graph.as_graph_def().node, 3)
     if name is not None:
       self._assert_all_graph_nodes_in_name_scope(graph, name)
@@ -178,7 +182,7 @@ def test_encode_decorator(self, name):
 
     # The graph should contain three nodes. The two above Python constants
     # converted to a Tensor object, and the resulting sum.
-    graph = tf.get_default_graph()
+    graph = tf.compat.v1.get_default_graph()
     self.assertLen(graph.as_graph_def().node, 3)
     if name is not None:
       self._assert_all_graph_nodes_in_name_scope(graph, name)
@@ -215,7 +219,7 @@ def test_decode_decorator(self, name):
 
     # The graph should contain six nodes. The four above Python constants
     # converted to a Tensor object, the subtraction, and the final reshape.
-    graph = tf.get_default_graph()
+    graph = tf.compat.v1.get_default_graph()
     self.assertLen(graph.as_graph_def().node, 6)
     if name is not None:
       self._assert_all_graph_nodes_in_name_scope(graph, name)
@@ -246,8 +250,8 @@ class NoneStateAdaptiveEncodingStageTest(tf.test.TestCase,
 
   def test_as_adaptive_encoding_stage(self):
     """Tests correctness of the wrapped encoding stage."""
-    a_var = tf.get_variable('a', initializer=2.0)
-    b_var = tf.get_variable('b', initializer=3.0)
+    a_var = tf.compat.v1.get_variable('a', initializer=2.0)
+    b_var = tf.compat.v1.get_variable('b', initializer=3.0)
     stage = test_utils.SimpleLinearEncodingStage(a_var, b_var)
     wrapped_stage = encoding_stage.as_adaptive_encoding_stage(stage)
     self.assertIsInstance(wrapped_stage,
@@ -279,7 +283,7 @@ def test_as_adaptive_encoding_stage(self):
     self.assertEqual(stage.decode_needs_input_shape,
                      wrapped_stage.decode_needs_input_shape)
 
-    self.evaluate(tf.global_variables_initializer())
+    self.evaluate(tf.compat.v1.global_variables_initializer())
     test_data = test_utils.TestData(*self.evaluate([x, encoded_x, decoded_x]))
     self.assertEqual(2.0, test_data.x)
     self.assertEqual(

tensorflow_model_optimization/python/core/internal/tensor_encoding/core/gather_encoder.py

@@ -390,7 +390,7 @@ def initial_state(self, name=None):
     Returns:
       A tuple of `Tensor` values, representing the initial state.
     """
-    with tf.name_scope(name, 'gather_encoder_initial_state'):
+    with tf.compat.v1.name_scope(name, 'gather_encoder_initial_state'):
       return self._initial_state_fn()
 
   def get_params(self, state=None, name=None):
@@ -417,7 +417,8 @@ def get_params(self, state=None, name=None):
     """
     if state is None:
       state = self.initial_state()
-    with tf.name_scope(name, 'gather_encoder_get_params', list(state)):
+    with tf.compat.v1.name_scope(name, 'gather_encoder_get_params',
+                                 list(state)):
       state = tf.nest.map_structure(tf.convert_to_tensor, state)
       return self._get_params_fn(state)
 
@@ -445,7 +446,7 @@ def encode(self, x, encode_params, name=None):
         `get_params` method.
     """
     values = [x] + list(encode_params)
-    with tf.name_scope(name, 'gather_encoder_encode', values):
+    with tf.compat.v1.name_scope(name, 'gather_encoder_encode', values):
       x = tf.convert_to_tensor(x)
       encode_params = tf.nest.map_structure(tf.convert_to_tensor, encode_params)
       return self._encode_fn(x, encode_params)
@@ -476,7 +477,8 @@ def decode_before_sum(self, encoded_x, decode_before_sum_params, name=None):
         `get_params` method.
     """
     values = list(encoded_x.values()) + list(decode_before_sum_params)
-    with tf.name_scope(name, 'gather_encoder_decode_before_sum', values):
+    with tf.compat.v1.name_scope(name, 'gather_encoder_decode_before_sum',
+                                 values):
       encoded_x = tf.nest.map_structure(tf.convert_to_tensor, encoded_x)
       decode_before_sum_params = tf.nest.map_structure(
           tf.convert_to_tensor, decode_before_sum_params)
@@ -513,7 +515,8 @@ def decode_after_sum(self,
     values = list(part_decoded_x.values()) if isinstance(
         part_decoded_x, dict) else [part_decoded_x]
     values = (values + list(decode_after_sum_params) + [num_summands])
-    with tf.name_scope(name, 'gather_encoder_decode_after_sum', values):
+    with tf.compat.v1.name_scope(name, 'gather_encoder_decode_after_sum',
+                                 values):
       part_decoded_x = tf.nest.map_structure(tf.convert_to_tensor,
                                              part_decoded_x)
       decode_after_sum_params = tf.nest.map_structure(tf.convert_to_tensor,
@@ -545,7 +548,7 @@ def update_state(self, state, state_update_tensors, name=None):
         return value of the `initial_state` method.
     """
     values = list(state) + list(state_update_tensors)
-    with tf.name_scope(name, 'gather_encoder_update_state', values):
+    with tf.compat.v1.name_scope(name, 'gather_encoder_update_state', values):
       state = tf.nest.map_structure(tf.convert_to_tensor, state)
       state_update_tensors = tf.nest.map_structure(tf.convert_to_tensor,
                                                    state_update_tensors)

tensorflow_model_optimization/python/core/internal/tensor_encoding/core/simple_encoder.py

@@ -161,7 +161,7 @@ def initial_state(self, name=None):
     Returns:
       A tuple of `Tensor` values, representing the initial state.
     """
-    with tf.name_scope(name, 'simple_encoder_initial_state'):
+    with tf.compat.v1.name_scope(name, 'simple_encoder_initial_state'):
       return self._initial_state_fn()
 
   def encode(self, x, state=None, name=None):
@@ -188,7 +188,8 @@ def encode(self, x, state=None, name=None):
     """
     if state is None:
       state = self.initial_state()
-    with tf.name_scope(name, 'simple_encoder_encode', [x] + list(state)):
+    with tf.compat.v1.name_scope(name, 'simple_encoder_encode',
+                                 [x] + list(state)):
       return self._encode_fn(x, state)
 
   def decode(self, encoded_x, name=None):
@@ -208,5 +209,6 @@ def decode(self, encoded_x, name=None):
         If `encoded_x` is not of the same structure as returned by the `encode`
         method.
     """
-    with tf.name_scope(name, 'simple_encoder_decode', encoded_x.values()):
+    with tf.compat.v1.name_scope(name, 'simple_encoder_decode',
+                                 encoded_x.values()):
       return self._decode_fn(encoded_x)

tensorflow_model_optimization/python/core/internal/tensor_encoding/stages/research/clipping_test.py

@@ -26,6 +26,10 @@
 from tensorflow_model_optimization.python.core.internal.tensor_encoding.testing import test_utils
 
 
+if tf.executing_eagerly():
+  tf.compat.v1.disable_eager_execution()
+
+
 class ClipByNormEncodingStageTest(test_utils.BaseEncodingStageTest):
 
   def default_encoding_stage(self):

tensorflow_model_optimization/python/core/internal/tensor_encoding/stages/research/kashin.py

@@ -206,7 +206,7 @@ def encode(self, x, encode_params):
       # If there is clipping in the last iteration, this can result in
       # biased representation of smaller magnitude. We compensate for this
       # by scaling such that the norm is preserved.
-      kashin_coefficients *= tf.div_no_nan(
+      kashin_coefficients *= tf.compat.v1.div_no_nan(
           tf.norm(x, axis=1, keepdims=True),
           tf.norm(kashin_coefficients, axis=1, keepdims=True))
 

tensorflow_model_optimization/python/core/internal/tensor_encoding/stages/research/kashin_test.py

@@ -26,6 +26,10 @@
 from tensorflow_model_optimization.python.core.internal.tensor_encoding.testing import test_utils
 
 
+if tf.executing_eagerly():
+  tf.compat.v1.disable_eager_execution()
+
+
 class KashinHadamardEncodingStageTest(test_utils.BaseEncodingStageTest):
 
   def default_encoding_stage(self):