In files where we explicitly import tf.compat.v1, don't qualify access of tf.compat.v1 again.

PiperOrigin-RevId: 289135863

Author: afrozenator

tensor2tensor/data_generators/generator_utils.py

@@ -867,18 +867,18 @@ def dict_pack(example):
 
   def _standardize(self, dataset, keys):
     """Force dataset structure into a tuple of Tensors."""
-    shapes = tf.compat.v1.data.get_output_shapes(dataset)
+    shapes = tf.data.get_output_shapes(dataset)
 
     if isinstance(shapes, dict):
       keys = keys or tuple(shapes.keys())
       dataset = dataset.map(lambda x: tuple(x[k] for k in keys))
-      shapes = tf.compat.v1.data.get_output_shapes(dataset)
+      shapes = tf.data.get_output_shapes(dataset)
 
     if not all(isinstance(i, tf.TensorShape) for i in shapes):
       # Internally this class expects tuples of Tensors, even for the degenerate
       # case of a single sequence.
       dataset = dataset.map(lambda x: (x,))
-      shapes = tf.compat.v1.data.get_output_shapes(dataset)
+      shapes = tf.data.get_output_shapes(dataset)
 
     for s in shapes:
       if not s.is_compatible_with(tf.TensorShape([None])):
@@ -890,7 +890,7 @@ def _standardize(self, dataset, keys):
     if self._chop_long_sequences and len(shapes) != 1:
       raise ValueError("chop_long_sequences expects a single sequence dataset.")
 
-    token_types = tf.compat.v1.data.get_output_types(dataset)
+    token_types = tf.data.get_output_types(dataset)
     if len(set(token_types)) > 1:
       raise ValueError("Inconsistent dtypes: {}".format(token_types))
 

tensor2tensor/data_generators/problem_test.py

@@ -30,7 +30,7 @@
 from tensor2tensor.utils import test_utils
 
 import tensorflow.compat.v1 as tf
-tf.compat.v1.enable_eager_execution()
+tf.enable_eager_execution()
 
 
 def assert_tensors_equal(sess, t1, t2, n):

tensor2tensor/layers/common_video_test.py

@@ -26,7 +26,7 @@
 from tensor2tensor.utils import test_utils
 
 import tensorflow.compat.v1 as tf
-tf.compat.v1.enable_eager_execution()
+tf.enable_eager_execution()
 
 
 class CommonVideoTest(parameterized.TestCase, tf.test.TestCase):

tensor2tensor/layers/discretization_test.py

@@ -24,7 +24,7 @@
 from tensor2tensor.utils import test_utils
 
 import tensorflow.compat.v1 as tf
-tf.compat.v1.enable_eager_execution()
+tf.enable_eager_execution()
 
 
 class DiscretizationTest(tf.test.TestCase):

tensor2tensor/layers/latent_layers_test.py

@@ -28,7 +28,7 @@
 from tensor2tensor.utils import test_utils
 
 import tensorflow.compat.v1 as tf
-tf.compat.v1.enable_eager_execution()
+tf.enable_eager_execution()
 
 
 def imagetransformer_latent_tiny():

tensor2tensor/layers/modalities_test.py

@@ -26,7 +26,7 @@
 from tensor2tensor.utils import test_utils
 
 import tensorflow.compat.v1 as tf
-tf.compat.v1.enable_eager_execution()
+tf.enable_eager_execution()
 
 
 class ModalityTest(tf.test.TestCase):

tensor2tensor/layers/ngram_test.py

@@ -24,7 +24,7 @@
 from tensor2tensor.utils import test_utils
 
 import tensorflow.compat.v1 as tf
-tf.compat.v1.enable_eager_execution()
+tf.enable_eager_execution()
 
 
 class NGramTest(tf.test.TestCase):

tensor2tensor/utils/adafactor.py

@@ -24,7 +24,7 @@
 import tensorflow.compat.v1 as tf
 
 
-class AdafactorOptimizer(tf.compat.v1.train.Optimizer):
+class AdafactorOptimizer(tf.train.Optimizer):
   """Optimizer that implements the Adafactor algorithm.
 
   Adafactor is described in https://arxiv.org/abs/1804.04235.

tensor2tensor/utils/multistep_optimizer.py

@@ -29,7 +29,7 @@
 import tensorflow.compat.v1 as tf
 
 
-class MultistepAdamOptimizer(tf.compat.v1.train.AdamOptimizer):
+class MultistepAdamOptimizer(tf.train.AdamOptimizer):
   """Adam with SGD updates every n steps with accumulated gradients."""
 
   def __init__(self, learning_rate=0.001, beta1=0.9, beta2=0.999, epsilon=1e-8,

tensor2tensor/utils/optimize.py

@@ -184,7 +184,7 @@ def _register_base_optimizer(name, opt):
   _register_base_optimizer(_name, _opt)
 
 
-class ConditionalOptimizer(tf.compat.v1.train.Optimizer):
+class ConditionalOptimizer(tf.train.Optimizer):
   """Conditional optimizer."""
 
   def __init__(self, optimizer_name, lr, hparams, use_tpu=False):  # pylint: disable=super-init-not-called