def decorator(self, *args, **kwargs):
    # TODO(b/117110239): Re-enable.
    # with context.execution_mode(context.ASYNC):
    #   f(self, *args, **kwargs)

    with context.execution_mode(context.SYNC):
      f(self, *args, **kwargs)

 def testEagerIteratorAsync(self):
   with context.eager_mode(), context.execution_mode(context.ASYNC):
     val = 0
     dataset = dataset_ops.Dataset.range(10)
     for foo in dataset:
       self.assertEqual(val, foo.numpy())
       val += 1

 def testDatasetEagerIteration(self, execution_mode):
   with context.eager_mode(), context.execution_mode(execution_mode):
     val = 0
     dataset = dataset_ops.Dataset.range(10)
     for foo in dataset:
       self.assertEqual(val, foo.numpy())
       val += 1

  def _next_internal(self):
    """Returns a nested structure of `tf.Tensor`s containing the next element.
    """
    if not context.executing_eagerly():
      with ops.device(self._device):
        ret = gen_dataset_ops.iterator_get_next(
            self._iterator_resource,
            output_types=self._flat_output_types,
            output_shapes=self._flat_output_shapes)
      return self._structure._from_compatible_tensor_list(ret)  # pylint: disable=protected-access

    # This runs in sync mode as iterators use an error status to communicate
    # that there is no more data to iterate over.
    # TODO(b/77291417): Fix
    with context.execution_mode(context.SYNC):
      with ops.device(self._device):
        # TODO(ashankar): Consider removing this ops.device() contextmanager
        # and instead mimic ops placement in graphs: Operations on resource
        # handles execute on the same device as where the resource is placed.
        # NOTE(mrry): Here we use the "_sync" variant of `iterator_get_next`
        # because in eager mode this code will run synchronously on the calling
        # thread. Therefore we do not need to make a defensive context switch
        # to a background thread, and can achieve a small constant performance
        # boost by invoking the iterator synchronously.
        ret = gen_dataset_ops.iterator_get_next_sync(
            self._iterator_resource,
            output_types=self._flat_output_types,
            output_shapes=self._flat_output_shapes)

      return self._structure._from_compatible_tensor_list(ret)  # pylint: disable=protected-access

 def _next_internal(self):
   """Returns a nested structure of `tf.Tensor`s containing the next element.
   """
   # This runs in sync mode as iterators use an error status to communicate
   # that there is no more data to iterate over.
   # TODO(b/77291417): Fix
   with context.execution_mode(context.SYNC):
     return super(Iterator, self)._next_internal()

 def _next_internal(self):
   """Returns a nested structure of `tf.Tensor`s containing the next element.
   """
   # This runs in sync mode as iterators use an error status to communicate
   # that there is no more data to iterate over.
   # TODO(b/77291417): Fix
   with context.execution_mode(context.SYNC):
     with ops.device(self._device):
       flat_ret = ged_ops.experimental_function_buffering_resource_get_next(
           function_buffer_resource=self._buffering_resource,
           output_types=self._flat_output_types)
     return self._element_structure._from_tensor_list(flat_ret)

  def testCopyBetweenDevicesAsync(self):
    with context.execution_mode(context.ASYNC):
      x = constant_op.constant([[1., 2.], [3., 4.]])
      x = x.cpu()
      x = x.gpu()
      x = x.gpu()
      x = x.cpu()
      context.async_wait()

    # Invalid device
    with self.assertRaises(RuntimeError):
      x.gpu(context.context().num_gpus() + 1)
      context.async_wait()
    context.async_clear_error()

 def _run(self, func, num_iters, execution_mode=None):
   # call func to maybe warm up the GPU
   ctx = context.context()
   with context.execution_mode(execution_mode):
     func()
     if execution_mode == context.ASYNC:
       ctx.async_wait()
     start = time.time()
     for _ in xrange(num_iters):
       func()
     if execution_mode == context.ASYNC:
       ctx.async_wait()
     end = time.time()
     mean_us = (end - start) * 1e6 / num_iters
     self.report_benchmark(
         iters=num_iters,
         wall_time=mean_us,
         extras={"examples_per_sec": num_iters / (end - start)})

 def testExecuteBasicAsync(self):
   with context.execution_mode(context.ASYNC):
     three = constant_op.constant(3)
     five = constant_op.constant(5)
     product = execute(
         b'Mul',
         num_outputs=1,
         inputs=[three, five],
         attrs=('T', three.dtype.as_datatype_enum))[0]
     self.assertAllEqual(15, product)
   # Error: Invalid arguments
   context.set_execution_mode(context.ASYNC)
   with self.assertRaises(errors.InvalidArgumentError):
     execute(
         b'MatMul',
         num_outputs=1,
         inputs=[three, five],
         attrs=('transpose_a', False, 'transpose_b', False, 'T',
                three.dtype.as_datatype_enum))
     context.async_wait()
   context.async_clear_error()
   context.set_execution_mode(context.SYNC)

  def decorator(self, *args, **kwargs):
    with context.execution_mode(context.ASYNC):
      f(self, *args, **kwargs)

    with context.execution_mode(context.SYNC):
      f(self, *args, **kwargs)