def set_optimizer_experimental_options(options):
  """Set experimental optimizer options.

  Note that optimizations are only applied in graph mode, (within tf.function).
  In addition, as these are experimental options, the list is subject to change.

  Args:
    options: Dictionary of experimental optimizer options to configure.
      Valid keys:
      - layout_optimizer: Optimize tensor layouts
        e.g. This will try to use NCHW layout on GPU which is faster.
      - constant_folding: Fold constants
        Statically infer the value of tensors when possible, and materialize the
        result using constants.
      - shape_optimization: Simplify computations made on shapes.
      - remapping: Remap subgraphs onto more efficient implementations.
      - arithmetic_optimization: Simplify arithmetic ops with common
        sub-expression elimination and arithmetic simplification.
      - dependency_optimization: Control dependency optimizations. Remove
        redundant control dependencies, which may enable other optimization.
        This optimizer is also essential for pruning Identity and NoOp nodes.
      - loop_optimization: Loop optimizations.
      - function_optimization: Function optimizations and inlining.
      - debug_stripper: Strips debug-related nodes from the graph.
      - disable_model_pruning: Disable removal of unnecessary ops from the graph
      - scoped_allocator_optimization: Try to allocate some independent Op
        outputs contiguously in order to merge or eliminate downstream Ops.
      - pin_to_host_optimization: Force small ops onto the CPU.
      - implementation_selector: Enable the swap of kernel implementations based
        on the device placement.
      - disable_meta_optimizer: Disable the entire meta optimizer.
      - min_graph_nodes: The minimum number of nodes in a graph to optimizer.
        For smaller graphs, optimization is skipped.
  """
  context.context().set_optimizer_experimental_options(options)

def set_visible_devices(devices, device_type=None):
  """Set the list of visible devices.

  Sets the list of PhysicalDevices to be marked as visible to the runtime. Any
  devices that are not marked as visible means TensorFlow will not allocate
  memory on it and will not be able to place any operations on it as no
  LogicalDevice will be created on it. By default all discovered devices are
  marked as visible.

  The following example demonstrates disabling the first GPU on the machine.

  ```python
  physical_devices = config.experimental.list_physical_devices('GPU')
  assert len(physical_devices) > 0, "Not enough GPU hardware devices available"
  # Disable first GPU
  tf.config.experimental.set_visible_devices(physical_devices[1:], 'GPU')
  logical_devices = config.experimental.list_logical_devices('GPU')
  # Logical device was not created for first GPU
  assert len(logical_devices) == len(physical_devices) - 1
  ```

  Args:
    devices: (optional) List of PhysicalDevice objects to make visible
    device_type: (optional) Device types to limit visibility configuration to.
      Other device types will be left unaltered.
  """
  context.context().set_visible_devices(devices, device_type)

 def testBadConstructorArgs(self):
   ctx = context.context()
   handle = ctx._handle
   device = ctx.device_name
   # Missing context.
   with self.assertRaisesRegexp(
       TypeError, r"Required argument 'context' \(pos 2\) not found"):
     ops.EagerTensor(1, device=device)
   # Missing device.
   with self.assertRaisesRegexp(
       TypeError, r"Required argument 'device' \(pos 3\) not found"):
     ops.EagerTensor(1, context=handle)
   # Bad dtype type.
   with self.assertRaisesRegexp(TypeError,
                                "Expecting a DataType value for dtype. Got"):
     ops.EagerTensor(1, context=handle, device=device, dtype="1")
   # Following errors happen when trying to copy to GPU.
   if not context.context().num_gpus():
     self.skipTest("No GPUs found")
   with ops.device("/device:GPU:0"):
     device = ctx.device_name
     # Bad context.
     with self.assertRaisesRegexp(
         TypeError, "Expecting a PyCapsule encoded context handle. Got"):
       ops.EagerTensor(1.0, context=1, device=device)
     # Bad device.
     with self.assertRaisesRegexp(
         TypeError, "Error parsing device argument to CopyToDevice"):
       ops.EagerTensor(1.0, context=handle, device=1)

  def testGpuInvalidConfig(self):
    gpus = config.list_physical_devices('GPU')
    self.assertNotEqual(len(gpus), 0)

    for gpu in gpus:
      config.set_memory_growth(gpu, True)

    c = context.context().config
    self.assertTrue(c.gpu_options.allow_growth)

    with self.assertRaisesRegexp(ValueError, 'memory limit'):
      config.set_virtual_device_configuration(gpus[-1], [
          context.VirtualDeviceConfiguration(),
          context.VirtualDeviceConfiguration()
      ])

    self.assertIsNone(config.get_virtual_device_configuration(gpus[-1]))
    config.set_virtual_device_configuration(gpus[-1], [
        context.VirtualDeviceConfiguration(memory_limit=10),
        context.VirtualDeviceConfiguration(memory_limit=10)
    ])

    c = context.context().config
    self.assertFalse(c.gpu_options.allow_growth)

    with self.assertRaisesRegexp(ValueError, 'virtual devices'):
      config.set_memory_growth(gpus[-1], False)

  def __init__(self, dataset):
    """Creates a new iterator over the given dataset.

    For example:
    ```python
    dataset = tf.data.Dataset.range(4)
    for x in Iterator(dataset):
      print(x)
    ```

    Tensors produced will be placed on the device on which this iterator object
    was created.

    Args:
      dataset: A `tf.data.Dataset` object.

    Raises:
      TypeError: If `dataset` is an unsupported type.
      RuntimeError: When invoked without eager execution enabled.
    """
    if not context.context().device_spec.device_type:
      is_remote_device = False
    else:
      is_remote_device = context.context().device_spec.device_type != "CPU"
    if is_remote_device:
      with ops.device(None):
        # Let the placer figure out where to place the various functions etc.
        # created by the CopyToDeviceDataset.
        dataset = dataset.apply(prefetching_ops.copy_to_device(
            context.context().device_name))
        dataset = dataset.prefetch(1)
    super(Iterator, self).__init__(dataset)

  def testJit(self):
    self.assertEqual(config.get_optimizer_jit(), False)

    # the following function should cause Op fusion to occur. However, there is
    # unfortunately no straightforward way to ensure this. We will just have to
    # settle for creating a test that can trigger JIT.
    @def_function.function
    def fun(a, b):
      c = a * b
      d = c + a
      return d

    a = constant_op.constant([2., 2.])
    b = constant_op.constant([2., 2.])

    self.evaluate(fun(a, b))

    config.set_optimizer_jit(True)
    self.assertEqual(config.get_optimizer_jit(), True)
    self.assertEqual(config.get_optimizer_jit(),
                     context.context().optimizer_jit)

    self.evaluate(fun(a, b))

    config.set_optimizer_jit(False)
    self.assertEqual(config.get_optimizer_jit(), False)
    self.assertEqual(config.get_optimizer_jit(),
                     context.context().optimizer_jit)

    self.evaluate(fun(a, b))

  def testSoftPlacement(self):
    if context.executing_eagerly():
      self.assertTrue(config.get_soft_device_placement())
    else:
      self.assertFalse(config.get_soft_device_placement())

    @def_function.function
    def mod():
      with ops.device('/device:GPU:0'):
        a = constant_op.constant(1.0)
        b = constant_op.constant(1.0)
        return math_ops.mod(a, b)

    config.set_soft_device_placement(True)
    self.assertEqual(config.get_soft_device_placement(), True)
    self.assertEqual(
        config.get_soft_device_placement(),
        context.context().soft_device_placement)

    # Since soft placement is enabled, the mod operation should work with CPU
    mod()

    config.set_soft_device_placement(False)
    self.assertEqual(config.get_soft_device_placement(), False)
    self.assertEqual(
        config.get_soft_device_placement(),
        context.context().soft_device_placement)

    # Since soft placement is disabled, the mod operation should fail on GPU
    with self.assertRaises(errors.InvalidArgumentError):
      mod()

 def testCopyScope(self):
   if not context.context().num_gpus():
     self.skipTest('No GPUs found')
   constant = constant_op.constant(1.0)
   with ops.device('gpu:0'):
     with context.context().device_policy(context.DEVICE_PLACEMENT_SILENT):
       c = constant + 1.0
   self.assertAllEqual(c, 2.0)

 def __del__(self):
   if self._created_eagerly:
     try:
       context.context().end_step()
     except AttributeError:
       pass
     except TypeError:
       pass

  def __init__(self, dataset):
    """Creates a new iterator over the given dataset.

    For example:
    ```python
    dataset = tf.data.Dataset.range(4)
    for x in Iterator(dataset):
      print(x)
    ```

    Tensors produced will be placed on the device on which this iterator object
    was created.

    Args:
      dataset: A `tf.data.Dataset` object.

    Raises:
      TypeError: If `dataset` is an unsupported type.
      RuntimeError: When invoked without eager execution enabled.
    """
    if isinstance(dataset, prefetching_ops._PrefetchToDeviceDataset):  # pylint: disable=protected-access
      raise TypeError(
          "`tf.contrib.data.prefetch_to_device()` is not compatible with "
          "`tf.contrib.eager.Iterator`. Use `for ... in dataset:` to iterate "
          "over the dataset instead.")

    super(Iterator, self).__init__(dataset)
    if not context.context().device_spec.device_type:
      is_remote_device = False
    else:
      is_remote_device = context.context().device_spec.device_type != "CPU"
    self._buffer_resource_handle = None
    if is_remote_device:
      with ops.device("/device:CPU:0"):
        iter_string_handle = gen_dataset_ops.iterator_to_string_handle(
            self._resource)

        @function.Defun(dtypes.string)
        def remote_fn(h):
          remote_iterator = iterator_ops.Iterator.from_string_handle(
              h, self.output_types, self.output_shapes, self.output_classes)
          return remote_iterator.get_next()

        remote_fn.add_to_graph(None)
        target = constant_op.constant("/device:CPU:0")
      with ops.device(self._device):
        self._buffer_resource_handle = prefetching_ops.function_buffering_resource(  # pylint: disable=line-too-long
            string_arg=iter_string_handle,
            output_types=self._flat_output_types,
            f=remote_fn,
            target_device=target,
            buffer_size=10,
            container="",
            shared_name=_generate_shared_name(
                "contrib_eager_iterator_function_buffer_resource"))
        self._buffer_resource_deleter = resource_variable_ops.EagerResourceDeleter(  # pylint: disable=line-too-long
            handle=self._buffer_resource_handle,
            handle_device=self._device)

  def testV1CompatibilityDummyInivisibleDeviceList(self):
    gpus = config.list_physical_devices('GPU')
    if gpus:
      self.skipTest('Test requires no GPUs')

    # Ensure GPU options left untouched on CPU only environments
    context.context()._physical_devices = None
    context.context()._config = config_pb2.ConfigProto(
        gpu_options=config_pb2.GPUOptions(visible_device_list='0'))
    new_config = context.context().config
    self.assertEqual(new_config.gpu_options.visible_device_list, '0')

 def as_default(self):
   """Enables summary writing within a `with` block."""
   if self._resource is None:
     yield self
   else:
     old = context.context().summary_writer_resource
     context.context().summary_writer_resource = self._resource
     yield self
     # Flushes the summary writer in eager mode or in graph functions, but not
     # in legacy graph mode (you're on your own there).
     self.flush()
     context.context().summary_writer_resource = old

  def testBenchmarks(self):
    # This isn't actually a test, but benchmarks packaged as a test
    # so that continuous integration runs catch any breakages.
    print(context.context())
    benchmark_create_tensor(FLAGS.iters or 30000)
    benchmark_matmul([2, 2], FLAGS.iters or 30000)
    benchmark_matmul([100, 28 * 28], FLAGS.iters or 1000)

    if context.context().num_gpus() > 0:
      print("---- RUNNING ON GPU NOW ----")
      benchmark_matmul([2, 2], FLAGS.iters or 30000, use_gpu=True)
      benchmark_matmul([100, 28 * 28], FLAGS.iters or 1000, use_gpu=True)

 def as_default(self):
   if self._resource is None:
     yield
   else:
     old = context.context().summary_writer_resource
     context.context().summary_writer_resource = self._resource
     yield
     # Flushes the summary writer in eager mode or in graph functions, but not
     # in legacy graph mode (you're on your own there).
     with ops.device("cpu:0"):
       gen_summary_ops.flush_summary_writer(self._resource)
     context.context().summary_writer_resource = old

  def __init__(self, persistent=False):
    """Creates a new GradientTape.

    Args:
      persistent: Boolean controlling whether a persistent gradient tape
        is created. False by default, which means at most one call can
        be made to the gradient() method on this object.
    """
    self._tape = None
    self._persistent = persistent
    self._recording = False
    context.context().start_step()

  def testContextSwitchStackContainsEagerMode(self):
    # Eager execution has been enabled, and no other context switch has
    # occurred, so `context_switches` should contain exactly one entry.
    self.assertEqual(len(context.context().context_switches.stack), 1)
    switch = context.context().context_switches.stack[0]

    # The entry should log that eager mode was entered.
    self.assertIs(switch.enter_context_fn, context.eager_mode)

    # It is not possible to build a graph function when eager execution
    # is enabled; the stack entry should reflect this fact.
    self.assertFalse(switch.is_building_function)

  def _benchmark_tfe_py_execute_matmul(self, m, transpose_b, num_iters):
    inputs = [m, m]
    # pylint: disable=protected-access
    ctx_handle = context.context()._handle
    # pylint: enable=protected-access
    device = context.context().device_name
    attrs = ("transpose_a", False, "transpose_b", transpose_b, "T",
             m.dtype.as_datatype_enum)
    def func():
      pywrap_tensorflow.TFE_Py_Execute(ctx_handle, device, "MatMul",
                                       inputs, attrs, 1)

    self._run(func, num_iters)

  def testCopyBetweenDevices(self):
    if not context.context().num_gpus():
      self.skipTest('No GPUs found')

    x = constant_op.constant([[1., 2.], [3., 4.]])
    x = x.as_cpu_tensor()
    x = x.as_gpu_tensor()
    x = x.as_gpu_tensor()
    x = x.as_cpu_tensor()

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

 def as_default(self):
   """Returns a context manager that enables summary writing."""
   if self._v2 and context.executing_eagerly() and self._closed:
     raise RuntimeError("SummaryWriter is already closed")
   old = context.context().summary_writer
   try:
     context.context().summary_writer = self
     yield self
     # Flushes the summary writer in eager mode or in graph functions, but
     # not in legacy graph mode (you're on your own there).
     self.flush()
   finally:
     context.context().summary_writer = old

  def testCopyBetweenDevices(self):
    if not context.context().num_gpus():
      self.skipTest('No GPUs found')

    x = tensor.Tensor([[1., 2.], [3., 4.]])
    x = x.as_cpu_tensor()
    x = x.as_gpu_tensor()
    x = x.as_gpu_tensor()
    x = x.as_cpu_tensor()

    # Invalid device
    with self.assertRaises(errors.InvalidArgumentError):
      x.as_gpu_tensor(context.context().num_gpus() + 1)