text: # coding=utf-8
  # Copyright 2019 The Tensor2Tensor Authors.
  #
  # 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.
  
  """Test utilities."""
  
  from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  
  import tensorflow.compat.v1 as tf
  
  
  def run_in_graph_and_eager_modes(func=None,
                                   config=None,
                                   use_gpu=True):
    """Execute the decorated test with and without enabling eager execution.
  
    This function returns a decorator intended to be applied to test methods in
    a `tf.test.TestCase` class. Doing so will cause the contents of the test
    method to be executed twice - once in graph mode, and once with eager
    execution enabled. This allows unittests to confirm the equivalence between
    eager and graph execution.
  
    NOTE: This decorator can only be used when executing eagerly in the
    outer scope.
  
    For example, consider the following unittest:
  
    ```python
    tf.enable_eager_execution()
  
    class SomeTest(tf.test.TestCase):
  
      @test_utils.run_in_graph_and_eager_modes
      def test_foo(self):
        x = tf.constant([1, 2])
        y = tf.constant([3, 4])
        z = tf.add(x, y)
        self.assertAllEqual([4, 6], self.evaluate(z))
  
    if __name__ == "__main__":
      tf.test.main()
    ```
  
    This test validates that `tf.add()` has the same behavior when computed with
    eager execution enabled as it does when constructing a TensorFlow graph and
    executing the `z` tensor with a session.
  
    Args:
      func: function to be annotated. If `func` is None, this method returns a
        decorator the can be applied to a function. If `func` is not None this
        returns the decorator applied to `func`.
      config: An optional config_pb2.ConfigProto to use to configure the session
        when executing graphs.
      use_gpu: If True, attempt to run as many operations as possible on GPU.
  
    Returns:
      Returns a decorator that will run the decorated test method twice:
      once by constructing and executing a graph in a session and once with
      eager execution enabled.
    """
  
    def decorator(f):
      """Decorator for a method."""
      def decorated(self, *args, **kwargs):
        """Run the decorated test method."""
        if not tf.executing_eagerly():
          raise ValueError("Must be executing eagerly when using the "
                           "run_in_graph_and_eager_modes decorator.")
  
        # Run eager block
        f(self, *args, **kwargs)
        self.tearDown()
  
        # Run in graph mode block
        with tf.Graph().as_default():
          self.setUp()
          with self.test_session(use_gpu=use_gpu, config=config):
            f(self, *args, **kwargs)
  
      return decorated
  
    if func is not None:
      return decorator(func)
  
    return decorator
  
  
  def run_in_graph_mode_only(func=None, config=None, use_gpu=True):
    """Runs a test in graph mode only, when eager is enabled by default."""
    def decorator(f):
      """Decorator for a method."""
      def decorated(self, *args, **kwargs):
        """Run the decorated test method."""
        self.tearDown()
        # Run in graph mode block
        with tf.Graph().as_default():
          self.setUp()
          with self.test_session(use_gpu=use_gpu, config=config):
            f(self, *args, **kwargs)
  
      return decorated
  
    if func is not None:
      return decorator(func)
  
    return decorator
  
  
  def test_main():
    tf.enable_eager_execution()
    tf.test.main()
  