Avoid using private APIs for determining if eager execution is enabled.

PiperOrigin-RevId: 193681604

Author: T2T Team

tensor2tensor/data_generators/image_utils.py

@@ -31,8 +31,6 @@
 
 import tensorflow as tf
 
-from tensorflow.python.eager import context
-
 
 def resize_by_area(img, size):
   """image resize function used by quite a few image problems."""
@@ -159,7 +157,7 @@ def generate_data(self, data_dir, tmp_dir, task_id=-1):
 
 
 def encode_images_as_png(images):
-  if context.in_eager_mode():
+  if tf.contrib.eager.in_eager_mode():
     for image in images:
       yield tf.image.encode_png(image).numpy()
   else:

tensor2tensor/layers/common_layers.py

@@ -32,7 +32,6 @@
 
 import tensorflow as tf
 
-from tensorflow.python.eager import context as tfe_context
 from tensorflow.python.framework import function
 from tensorflow.python.framework import ops
 
@@ -265,7 +264,7 @@ def embedding(x,
     # On the backwards pass, we want to convert the gradient from
     # an indexed-slices to a regular tensor before sending it back to the
     # parameter server. This avoids excess computation on the parameter server.
-    if not tfe_context.in_eager_mode():
+    if not tf.contrib.eager.in_eager_mode():
       embedding_var = eu.convert_gradient_to_tensor(embedding_var)
     x = dropout_no_scaling(x, 1.0 - symbol_dropout_rate)
     emb_x = gather(embedding_var, x, dtype)
@@ -2541,7 +2540,7 @@ def ones_matrix_band_part(rows, cols, num_lower, num_upper, out_shape=None):
 def reshape_like_all_dims(a, b):
   """Reshapes a to match the shape of b."""
   ret = tf.reshape(a, tf.shape(b))
-  if not tfe_context.in_eager_mode():
+  if not tf.contrib.eager.in_eager_mode():
     ret.set_shape(b.get_shape())
   return ret
 

tensor2tensor/layers/modalities.py

@@ -29,8 +29,6 @@
 
 import tensorflow as tf
 
-from tensorflow.python.eager import context
-
 
 @registry.register_symbol_modality("default")
 class SymbolModality(modality.Modality):
@@ -97,7 +95,7 @@ def _get_weights(self, hidden_dim=None):
     else:
       ret = tf.concat(shards, 0)
     # Convert ret to tensor.
-    if not context.in_eager_mode():
+    if not tf.contrib.eager.in_eager_mode():
       ret = eu.convert_gradient_to_tensor(ret)
     return ret
 
@@ -211,13 +209,13 @@ class ImageModality(modality.Modality):
   def bottom(self, inputs):
     with tf.variable_scope(self.name):
       inputs = tf.to_float(inputs)
-      if not context.in_eager_mode():
+      if not tf.contrib.eager.in_eager_mode():
         tf.summary.image("inputs", inputs, max_outputs=2)
       return inputs
 
   def targets_bottom(self, inputs):
     with tf.variable_scope(self.name):
-      if not context.in_eager_mode():
+      if not tf.contrib.eager.in_eager_mode():
         tf.summary.image("targets_bottom",
                          tf.cast(inputs, tf.uint8), max_outputs=1)
       inputs_shape = common_layers.shape_list(inputs)
@@ -466,7 +464,7 @@ def bottom(self, inputs):
         raise ValueError("Assuming videos given as tensors in the format "
                          "[batch, time, height, width, channels] but got one "
                          "of shape: %s" % str(inputs_shape))
-      if not context.in_eager_mode():
+      if not tf.contrib.eager.in_eager_mode():
         tf.summary.image("inputs", tf.cast(inputs[:, -1, :, :, :], tf.uint8),
                          max_outputs=1)
       # Standardize frames.
@@ -487,7 +485,7 @@ def targets_bottom(self, inputs):
         raise ValueError("Assuming videos given as tensors in the format "
                          "[batch, time, height, width, channels] but got one "
                          "of shape: %s" % str(inputs_shape))
-      if not context.in_eager_mode():
+      if not tf.contrib.eager.in_eager_mode():
         tf.summary.image(
             "targets_bottom", tf.cast(inputs[:, -1, :, :, :], tf.uint8),
             max_outputs=1)

tensor2tensor/utils/expert_utils.py

@@ -33,7 +33,6 @@
 from six.moves import zip  # pylint: disable=redefined-builtin
 import tensorflow as tf
 
-from tensorflow.python.eager import context
 from tensorflow.python.framework import function
 
 DEFAULT_DEV_STRING = "existing_device"
@@ -560,7 +559,7 @@ def remove(self, x):
           x,
           indices=self.nonpad_ids,
       )
-      if not context.in_eager_mode():
+      if not tf.contrib.eager.in_eager_mode():
         # This is a hack but for some reason, gather_nd return a tensor of
         # undefined shape, so the shape is set up manually
         x.set_shape([None] + x_shape[1:])
@@ -909,15 +908,15 @@ def my_fn(x):
 def reshape_like(a, b):
   """Reshapes a to match the shape of b in all but the last dimension."""
   ret = tf.reshape(a, tf.concat([tf.shape(b)[:-1], tf.shape(a)[-1:]], 0))
-  if not context.in_eager_mode():
+  if not tf.contrib.eager.in_eager_mode():
     ret.set_shape(b.get_shape().as_list()[:-1] + a.get_shape().as_list()[-1:])
   return ret
 
 
 def flatten_all_but_last(a):
   """Flatten all dimensions of a except the last."""
   ret = tf.reshape(a, [-1, tf.shape(a)[-1]])
-  if not context.in_eager_mode():
+  if not tf.contrib.eager.in_eager_mode():
     ret.set_shape([None] + a.get_shape().as_list()[-1:])
   return ret
 

tensor2tensor/utils/t2t_model.py

@@ -44,7 +44,6 @@
 
 import tensorflow as tf
 
-from tensorflow.python.eager import context
 from tensorflow.python.layers import base
 from tensorflow.python.ops import variable_scope
 
@@ -717,7 +716,7 @@ def _slow_greedy_infer(self, features, decode_length):
 
     def infer_step(recent_output, recent_logits, unused_loss):
       """Inference step."""
-      if not context.in_eager_mode():
+      if not tf.contrib.eager.in_eager_mode():
         recent_output.set_shape([None, None, None, 1])
       padded = tf.pad(recent_output, [[0, 0], [0, 1], [0, 0], [0, 0]])
       features["targets"] = padded
@@ -733,7 +732,7 @@ def infer_step(recent_output, recent_logits, unused_loss):
                              common_layers.shape_list(recent_output)[1], :, :]
       cur_sample = tf.to_int64(tf.expand_dims(cur_sample, axis=1))
       samples = tf.concat([recent_output, cur_sample], axis=1)
-      if not context.in_eager_mode():
+      if not tf.contrib.eager.in_eager_mode():
         samples.set_shape([None, None, None, 1])
 
       # Assuming we have one shard for logits.
@@ -765,7 +764,7 @@ def infer_step(recent_output, recent_logits, unused_loss):
     result = initial_output
     # tensor of shape [batch_size, time, 1, 1, vocab_size]
     logits = tf.zeros((batch_size, 0, 1, 1, target_modality.top_dimensionality))
-    if not context.in_eager_mode():
+    if not tf.contrib.eager.in_eager_mode():
       logits.set_shape([None, None, None, None, None])
     loss = 0.0
 
@@ -1304,7 +1303,7 @@ def as_default(self):
 
 
 def create_eager_var_store():
-  if context.in_eager_mode():
+  if tf.contrib.eager.in_eager_mode():
     return variable_scope.EagerVariableStore()
   else:
     return DummyVariableStore()
@@ -1405,7 +1404,7 @@ def summarize_features(features, num_shards=1):
 
 
 def _eager_log(level, *args):
-  if context.in_eager_mode() and args in _already_logged:
+  if tf.contrib.eager.in_eager_mode() and args in _already_logged:
     return
   _already_logged.add(args)
   getattr(tf.logging, level)(*args)