def testInsideFunction(self):
    if test_util.is_gpu_available():
      self.skipTest(
          "b/123899495: Colocation errors for critical sections in map on GPU")
    cs = critical_section_ops.CriticalSection()
    with ops.device("/gpu:0" if test_util.is_gpu_available() else "/cpu:0"):
      v = resource_variable_ops.ResourceVariable(1)
    def fn():
      return v.read_value()

    # map() creates a TensorFlow function.
    ds = dataset_ops.Dataset.range(1)
    if test_util.is_gpu_available():
      ds = (ds.apply(prefetching_ops.copy_to_device("/gpu:0"))
            .apply(prefetching_ops.map_on_gpu(lambda _: cs.execute(fn))))
    else:
      ds = ds.map(lambda _: cs.execute(fn))

    def get_first():
      if context.executing_eagerly():
        return self.evaluate(ds.make_one_shot_iterator().get_next())
      itr = ds.make_initializable_iterator()
      self.evaluate([v.initializer, itr.initializer])
      return self.evaluate(itr.get_next())

    self.assertEqual(1, get_first())

  def testCondAndTensorArrayInDefun(self):
    if test_util.is_gpu_available():
      old_enable_tensor_array_v2 = tensor_array_ops.ENABLE_TENSOR_ARRAY_V2
      # TODO(b/119689663): Enable this.
      tensor_array_ops.ENABLE_TENSOR_ARRAY_V2 = False

    @function.defun
    def f():
      x = math_ops.range(-5, 5)
      output = tensor_array_ops.TensorArray(dtype=dtypes.int32, size=x.shape[0])

      def loop_body(i, output):

        def if_true():
          return output.write(i, x[i]**2)

        def if_false():
          return output.write(i, x[i])

        output = control_flow_ops.cond(x[i] > 0, if_true, if_false)
        return i + 1, output

      _, output = control_flow_ops.while_loop(
          lambda i, arr: i < x.shape[0],
          loop_body,
          loop_vars=(constant_op.constant(0), output))
      return output.stack()

    output_t = f()
    self.assertAllEqual(
        self.evaluate(output_t), [-5, -4, -3, -2, -1, 0, 1, 4, 9, 16])

    if test_util.is_gpu_available():
      tensor_array_ops.ENABLE_TENSOR_ARRAY_V2 = old_enable_tensor_array_v2

def _generate_synthetic_snli_data_batch(sequence_length,
                                        batch_size,
                                        vocab_size):
  """Generate a fake batch of SNLI data for testing."""
  with tf.device("cpu:0"):
    labels = tf.random_uniform([batch_size], minval=1, maxval=4, dtype=tf.int64)
    prem = tf.random_uniform(
        (sequence_length, batch_size), maxval=vocab_size, dtype=tf.int64)
    prem_trans = tf.constant(np.array(
        [[3, 3, 2, 3, 3, 3, 2, 2, 2, 3, 3, 3,
          2, 3, 3, 2, 2, 3, 3, 3, 2, 2, 2, 2,
          3, 2, 2]] * batch_size, dtype=np.int64).T)
    hypo = tf.random_uniform(
        (sequence_length, batch_size), maxval=vocab_size, dtype=tf.int64)
    hypo_trans = tf.constant(np.array(
        [[3, 3, 2, 3, 3, 3, 2, 2, 2, 3, 3, 3,
          2, 3, 3, 2, 2, 3, 3, 3, 2, 2, 2, 2,
          3, 2, 2]] * batch_size, dtype=np.int64).T)
  if test_util.is_gpu_available():
    labels = labels.gpu()
    prem = prem.gpu()
    prem_trans = prem_trans.gpu()
    hypo = hypo.gpu()
    hypo_trans = hypo_trans.gpu()
  return labels, prem, prem_trans, hypo, hypo_trans

  def testUsingNamesNotUsingIntermediateTensors(self):
    if test_util.is_gpu_available():
      self.skipTest("b/123446705 this causes a segfault on GPU")

    with NodeStepper(self.sess, "e:0") as stepper:
      # The first cont() call should have used no feeds.
      result = stepper.cont("c:0")
      self.assertAllClose(6.0, result)
      self.assertItemsEqual(["a/read:0", "b/read:0"],
                            stepper.intermediate_tensor_names())
      self.assertAllClose(2.0, stepper.get_tensor_value("a/read:0"))
      self.assertAllClose(3.0, stepper.get_tensor_value("b/read:0"))
      self.assertEqual({}, stepper.last_feed_types())

      # The second cont() call should have used the tensor handle from the
      # previous cont() call.
      result = stepper.cont("e:0")
      self.assertAllClose(24.0, result)
      self.assertItemsEqual(["a/read:0", "b/read:0", "d:0"],
                            stepper.intermediate_tensor_names())
      self.assertAllClose(2.0, stepper.get_tensor_value("a/read:0"))
      self.assertAllClose(3.0, stepper.get_tensor_value("b/read:0"))
      self.assertAllClose(4.0, stepper.get_tensor_value("d:0"))
      self.assertEqual({
          "c:0": NodeStepper.FEED_TYPE_HANDLE,
          "a/read:0": NodeStepper.FEED_TYPE_DUMPED_INTERMEDIATE,
      }, stepper.last_feed_types())

  def test_function_with_captured_dataset(self):
    if test_util.is_gpu_available():
      self.skipTest("Currently broken when a GPU is available.")

    class HasDataset(module.Module):

      def __init__(self):
        super(HasDataset, self).__init__()
        self.dataset = (
            dataset_ops.Dataset.range(5)
            .map(lambda x: x ** 2))

      @def_function.function
      def __call__(self, x):
        current_sum = array_ops.zeros([], dtype=dtypes.int64)
        for element in self.dataset:
          current_sum += x * element
        return current_sum

    root = HasDataset()
    save_dir = os.path.join(self.get_temp_dir(), "saved_model")
    save.save(
        root, save_dir,
        signatures=root.__call__.get_concrete_function(
            tensor_spec.TensorSpec(None, dtypes.int64)))
    self.assertAllClose({"output_0": 3 * (1 + 4 + 9 + 16)},
                        _import_and_infer(save_dir, {"x": 3}))

  def testCopyToGPU(self):
    if not test_util.is_gpu_available():
      self.skipTest("No GPU available")

    with ops.device("/cpu:0"):
      optional_with_value = optional_ops.Optional.from_value(
          (constant_op.constant(37.0), constant_op.constant("Foo"),
           constant_op.constant(42)))
      optional_none = optional_ops.Optional.none_from_structure(
          structure.TensorStructure(dtypes.float32, []))

    with ops.device("/gpu:0"):
      gpu_optional_with_value = optional_ops._OptionalImpl(
          array_ops.identity(optional_with_value._variant_tensor),
          optional_with_value.value_structure)
      gpu_optional_none = optional_ops._OptionalImpl(
          array_ops.identity(optional_none._variant_tensor),
          optional_none.value_structure)

      gpu_optional_with_value_has_value = gpu_optional_with_value.has_value()
      gpu_optional_with_value_values = gpu_optional_with_value.get_value()

      gpu_optional_none_has_value = gpu_optional_none.has_value()

    self.assertTrue(self.evaluate(gpu_optional_with_value_has_value))
    self.assertEqual((37.0, b"Foo", 42),
                     self.evaluate(gpu_optional_with_value_values))
    self.assertFalse(self.evaluate(gpu_optional_none_has_value))

  def testCopyToDeviceGpuWithMap(self):
    if not test_util.is_gpu_available():
      self.skipTest("No GPU available")

    def generator():
      for i in range(10):
        yield i, float(i), str(i)

    host_dataset = dataset_ops.Dataset.from_generator(
        generator, output_types=(dtypes.int32, dtypes.float32, dtypes.string))
    device_dataset = host_dataset.apply(
        prefetching_ops.copy_to_device("/gpu:0"))

    def gpu_map_func(x, y, z):
      return math_ops.square(x), math_ops.square(y), z

    device_dataset = device_dataset.apply(
        prefetching_ops.map_on_gpu(gpu_map_func))
    options = dataset_ops.Options()
    options.experimental_autotune = False
    device_dataset = device_dataset.with_options(options)

    with ops.device("/gpu:0"):
      iterator = device_dataset.make_initializable_iterator()
      next_element = iterator.get_next()

    with self.cached_session() as sess:
      sess.run(iterator.initializer)
      for i in range(10):
        x, y, z = sess.run(next_element)
        self.assertEqual(i**2, x)
        self.assertEqual(float(i**2), y)
        self.assertEqual(util_compat.as_bytes(str(i)), z)
      with self.assertRaises(errors.OutOfRangeError):
        sess.run(next_element)

  def testAddN(self):
    devices = ["/cpu:0"]
    if test_util.is_gpu_available():
      devices.append("/gpu:0")
    for device in devices:
      with ops.device(device):
        # With value
        opt1 = optional_ops.Optional.from_value((1.0, 2.0))
        opt2 = optional_ops.Optional.from_value((3.0, 4.0))

        add_tensor = math_ops.add_n([opt1._variant_tensor,
                                     opt2._variant_tensor])
        add_opt = optional_ops._OptionalImpl(add_tensor, opt1.value_structure)
        self.assertAllEqual(self.evaluate(add_opt.get_value()), (4.0, 6.0))

        # Without value
        opt_none1 = optional_ops.Optional.none_from_structure(
            opt1.value_structure)
        opt_none2 = optional_ops.Optional.none_from_structure(
            opt2.value_structure)
        add_tensor = math_ops.add_n([opt_none1._variant_tensor,
                                     opt_none2._variant_tensor])
        add_opt = optional_ops._OptionalImpl(add_tensor,
                                             opt_none1.value_structure)
        self.assertFalse(self.evaluate(add_opt.has_value()))

  def testContWithPlaceholders(self):
    if test_util.is_gpu_available():
      self.skipTest("b/123446705 this causes a segfault on GPU")

    with NodeStepper(
        self.sess,
        self.y,
        feed_dict={
            self.ph0: [[1.0, 2.0], [-3.0, 5.0]],
            self.ph1: [[-1.0], [0.5]]
        }) as stepper:
      self.assertEqual(4, len(stepper.sorted_nodes()))
      self.assertSetEqual({"ph0:0", "ph1:0", "x:0", "y:0"},
                          set(stepper.closure_elements()))

      result = stepper.cont(self.x)
      self.assertAllClose([[0.0], [5.5]], result)
      self.assertEqual({
          "ph0:0": NodeStepper.FEED_TYPE_CLIENT,
          "ph1:0": NodeStepper.FEED_TYPE_CLIENT,
      }, stepper.last_feed_types())

      self.assertEqual(["x:0"], stepper.handle_names())
      self.assertSetEqual({"x"}, stepper.handle_node_names())

      result = stepper.cont(self.y)
      self.assertAllClose([[-1.0], [6.0]], result)
      self.assertEqual({
          "x:0": NodeStepper.FEED_TYPE_HANDLE,
          "ph1:0": NodeStepper.FEED_TYPE_CLIENT,
      }, stepper.last_feed_types())

  def testDeviceBeforeCond(self):
    with ops.Graph().as_default() as g:
      with self.session(graph=g):

        def fn():
          self.assertEqual("", constant_op.constant(3.0).op.device)
          return test_ops.device_placement_op()

        with ops.device("/device:CPU:0"):
          self.assertIn(
              compat.as_bytes("CPU:0"),
              self.evaluate(cond_v2.cond_v2(constant_op.constant(True),
                                            fn, fn)))

        def fn2():
          self.assertEqual("", constant_op.constant(3.0).op.device)
          return test_ops.device_placement_op()

        if test_util.is_gpu_available():
          with ops.device("/device:GPU:0"):
            self.assertIn(
                compat.as_bytes("GPU:0"),
                self.evaluate(cond_v2.cond_v2(constant_op.constant(True),
                                              fn2, fn2)))
        else:
          self.skipTest("Test requires a GPU to check GPU device placement.")

 def _compareScalar(self, func, x, y, dtype):
   with self.test_session(force_gpu=test_util.is_gpu_available()):
     out = func(
         ops.convert_to_tensor(np.array([x]).astype(dtype)),
         ops.convert_to_tensor(np.array([y]).astype(dtype)))
     ret = self.evaluate(out)
   return ret[0]

  def testDifferentDeviceCPUGPU(self):
    if not test_util.is_gpu_available():
      self.skipTest("No GPU available")

    self._prefetch_fn_helper_one_shot("cpu_gpu",
                                      "/job:localhost/replica:0/task:0/cpu:0",
                                      "/job:localhost/replica:0/task:0/gpu:0")

  def testGetNextAsOptionalGpu(self):
    if not test_util.is_gpu_available() or context.executing_eagerly():
      self.skipTest("No GPU available")

    dataset = dataset_ops.Dataset.range(9)
    multi_device_iterator = multi_device_iterator_ops.MultiDeviceIterator(
        dataset, ["/cpu:1", "/gpu:0"])
    elem_on_1, elem_on_2 = multi_device_iterator.get_next_as_optional()
    elem_on_1_has_value_t = elem_on_1.has_value()
    elem_on_1_t = elem_on_1.get_value()
    elem_on_2_has_value_t = elem_on_2.has_value()
    elem_on_2_t = elem_on_2.get_value()

    config = config_pb2.ConfigProto(device_count={"CPU": 2, "GPU": 1})
    with self.test_session(config=config) as sess:
      self.evaluate(multi_device_iterator.initializer)
      for i in range(0, 8, 2):
        elem_on_1_has_value, elem_on_1_value = sess.run(
            [elem_on_1_has_value_t, elem_on_1_t])
        self.assertTrue(elem_on_1_has_value)
        self.assertEqual(i, elem_on_1_value)
        elem_on_2_has_value, elem_on_2_value = sess.run(
            [elem_on_2_has_value_t, elem_on_2_t])
        self.assertTrue(elem_on_2_has_value)
        self.assertEqual(i + 1, elem_on_2_value)
      elem_on_1_has_value, elem_on_1_value = sess.run(
          [elem_on_1_has_value_t, elem_on_1_t])
      self.assertTrue(elem_on_1_has_value)
      self.assertEqual(8, elem_on_1_value)
      self.assertFalse(self.evaluate(elem_on_1_has_value_t))
      self.assertFalse(self.evaluate(elem_on_2_has_value_t))
      with self.assertRaises(errors.InvalidArgumentError):
        self.evaluate(elem_on_1_t)
      with self.assertRaises(errors.InvalidArgumentError):
        self.evaluate(elem_on_2_t)

  def testIteratorGetNextAsOptionalOnGPU(self):
    if not test_util.is_gpu_available():
      self.skipTest("No GPU available")

    host_dataset = dataset_ops.Dataset.range(3)
    device_dataset = host_dataset.apply(
        prefetching_ops.copy_to_device("/gpu:0"))
    with ops.device("/gpu:0"):
      iterator = device_dataset.make_initializable_iterator()
      next_elem = iterator_ops.get_next_as_optional(iterator)
      elem_has_value_t = next_elem.has_value()
      elem_value_t = next_elem.get_value()

    with self.cached_session() as sess:
      # Before initializing the iterator, evaluating the optional fails with
      # a FailedPreconditionError.
      with self.assertRaises(errors.FailedPreconditionError):
        sess.run(elem_has_value_t)
      with self.assertRaises(errors.FailedPreconditionError):
        sess.run(elem_value_t)

      # For each element of the dataset, assert that the optional evaluates to
      # the expected value.
      sess.run(iterator.initializer)
      for i in range(3):
        elem_has_value, elem_value = sess.run([elem_has_value_t, elem_value_t])
        self.assertTrue(elem_has_value)
        self.assertEqual(i, elem_value)

      # After exhausting the iterator, `next_elem.has_value()` will evaluate to
      # false, and attempting to get the value will fail.
      for _ in range(2):
        self.assertFalse(sess.run(elem_has_value_t))
        with self.assertRaises(errors.InvalidArgumentError):
          sess.run(elem_value_t)

  def testBadConstructorArgs(self):
    context.ensure_initialized()
    ctx = context.context()
    handle = ctx._handle
    device = ctx.device_name
    # Missing context.
    with self.assertRaisesRegexp(
        TypeError, r".*argument 'context' \(pos 2\).*"):
      ops.EagerTensor(1, device=device)
    # Missing device.
    with self.assertRaisesRegexp(
        TypeError, r".*argument 'device' \(pos 3\).*"):
      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 test_util.is_gpu_available():
      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 testFunctionWithResourcesOnDifferentDevices(self):
    if not test_util.is_gpu_available():
      self.skipTest("No GPUs available.")

    with ops.device("/cpu:0"):
      v_cpu_zero = resource_variable_ops.ResourceVariable(
          [0.0, 1.0, 2.0], name="v_cpu_zero")

    with ops.device("/cpu:1"):
      v_cpu_one = resource_variable_ops.ResourceVariable(
          [0.0, 1.0, 2.0], name="v_cpu_one")

    with ops.device("/gpu:0"):
      v_gpu = resource_variable_ops.ResourceVariable(
          [0.0, 1.0, 2.0], name="v_gpu")

    def sum_gather():
      cpu_result = math_ops.reduce_sum(array_ops.gather(v_cpu_zero, [1, 2]))
      also_cpu_result = math_ops.reduce_sum(array_ops.gather(v_cpu_one, [1, 2]))
      gpu_result = math_ops.reduce_sum(array_ops.gather(v_gpu, [1, 2]))
      return cpu_result, also_cpu_result, gpu_result

    defined = function.Defun()(sum_gather)
    with self.test_session(
        config=config_pb2.ConfigProto(
            allow_soft_placement=False,
            log_device_placement=True,
            device_count={"CPU": 2})) as sess:
      self.evaluate(variables.global_variables_initializer())
      expected = self.evaluate(sum_gather())
      result = sess.run(
          functional_ops.partitioned_call(
              args=defined.captured_inputs, f=defined))
      self.assertAllEqual(expected, result)

  def testInstantError(self):
    if test_util.is_gpu_available():
      # TODO(nareshmodi): make this test better
      self.skipTest("Gather doesn't do index checking on GPUs")

    with self.assertRaisesRegexp(errors.InvalidArgumentError,
                                 r'indices = 7 is not in \[0, 3\)'):
      array_ops.gather([0, 1, 2], 7)

 def _compare(self, c, x, y, use_gpu):
   np_ans = np.where(c, x, y)
   with self.test_session(use_gpu=use_gpu,
                          force_gpu=use_gpu and test_util.is_gpu_available()):
     out = array_ops.where(c, x, y)
     tf_ans = self.evaluate(out)
   self.assertAllEqual(np_ans, tf_ans)
   self.assertShapeEqual(np_ans, out)

 def testNoShape(self):
   with self.test_session(force_gpu=test_util.is_gpu_available()):
     p = array_ops.placeholder_with_default([17], shape=None)
     a = array_ops.identity(p)
     self.assertAllEqual([17], a.eval())
     self.assertAllEqual([3, 37], a.eval(feed_dict={p: [3, 37]}))
     self.assertAllEqual(
         [[3, 3], [3, 3]], a.eval(feed_dict={p: [[3, 3], [3, 3]]}))

  def testPartialShape(self):
    with self.test_session(force_gpu=test_util.is_gpu_available()):
      p = array_ops.placeholder_with_default([1, 2, 3], shape=[None])
      a = array_ops.identity(p)
      self.assertAllEqual([1, 2, 3], a.eval())
      self.assertAllEqual([3, 37], a.eval(feed_dict={p: [3, 37]}))

      with self.assertRaises(ValueError):
        a.eval(feed_dict={p: [[2, 2], [2, 2]]})