def benchmarkBatchMatMulBroadcast(self):
    for (a_shape, b_shape) in self.shape_pairs:
      with compat.forward_compatibility_horizon(2019, 4, 26):
        with ops.Graph().as_default(), \
            session.Session(config=benchmark.benchmark_config()) as sess, \
            ops.device("/cpu:0"):
          matrix_a = variables.Variable(
              GetRandomNormalInput(a_shape, np.float32))
          matrix_b = variables.Variable(
              GetRandomNormalInput(b_shape, np.float32))
          variables.global_variables_initializer().run()

          # Use batch matmul op's internal broadcasting.
          self.run_op_benchmark(
              sess,
              math_ops.matmul(matrix_a, matrix_b),
              min_iters=50,
              name="batch_matmul_cpu_{}_{}".format(a_shape, b_shape))

          # Manually broadcast the input matrices using the broadcast_to op.
          broadcasted_batch_shape = array_ops.broadcast_static_shape(
              matrix_a.shape[:-2], matrix_b.shape[:-2])
          broadcasted_a_shape = broadcasted_batch_shape.concatenate(
              matrix_a.shape[-2:])
          broadcasted_b_shape = broadcasted_batch_shape.concatenate(
              matrix_b.shape[-2:])
          self.run_op_benchmark(
              sess,
              math_ops.matmul(
                  array_ops.broadcast_to(matrix_a, broadcasted_a_shape),
                  array_ops.broadcast_to(matrix_b, broadcasted_b_shape)),
              min_iters=50,
              name="batch_matmul_manual_broadcast_cpu_{}_{}".format(
                  a_shape, b_shape))

  def testNMS128From1024(self):
    with compat.forward_compatibility_horizon(2018, 8, 8):
      num_boxes = 1024
      boxes_np = np.random.normal(50, 10, (num_boxes, 4)).astype("f4")
      scores_np = np.random.normal(0.5, 0.1, (num_boxes,)).astype("f4")

      max_output_size = 128
      iou_threshold_np = np.array(0.5, dtype=np.float32)
      score_threshold_np = np.array(0.0, dtype=np.float32)

      with self.cached_session() as sess:
        boxes = array_ops.placeholder(boxes_np.dtype, shape=boxes_np.shape)
        scores = array_ops.placeholder(scores_np.dtype, shape=scores_np.shape)
        iou_threshold = array_ops.placeholder(iou_threshold_np.dtype,
                                              iou_threshold_np.shape)
        score_threshold = array_ops.placeholder(score_threshold_np.dtype,
                                                score_threshold_np.shape)
        with self.test_scope():
          selected_indices = image_ops.non_max_suppression_padded(
              boxes=boxes,
              scores=scores,
              max_output_size=max_output_size,
              iou_threshold=iou_threshold,
              score_threshold=score_threshold,
              pad_to_max_output_size=True)
        inputs_feed = {
            boxes: boxes_np,
            scores: scores_np,
            score_threshold: score_threshold_np,
            iou_threshold: iou_threshold_np
        }
        (indices_tf, _) = sess.run(selected_indices, feed_dict=inputs_feed)

        self.assertEqual(indices_tf.size, max_output_size)

 def testStaticRegexReplaceDelegation(self):
   with compat.forward_compatibility_horizon(2018, 10, 11):
     with self.test_session():
       input_vector = constant_op.constant("foo", dtypes.string)
       pattern = "[a-z]"
       replace = "."
       op = string_ops.regex_replace(input_vector, pattern, replace)
       self.assertTrue(op.name.startswith("StaticRegexReplace"))

 def test3DTensorAsInputNoReshape(self):
   with compat.forward_compatibility_horizon(2018, 8, 27):
     self._testSoftmax(
         np.array([[[1., 1., 1., 1.], [1., 2., 3., 4.]],
                   [[2., 3., 4., 5.], [6., 7., 8., 9.]],
                   [[5., 4., 3., 2.], [1., 2., 3., 4.]]]).astype(np.float32),
         use_gpu=False)
     self._testOverflow(use_gpu=False)

  def test_decorator(self):
    compatibility_date = self._compatibility_date()
    one_day_after = self._n_days_after(1)
    with compat.forward_compatibility_horizon(*one_day_after):
      self.assertTrue(compat.forward_compatible(*compatibility_date))
      self.assertFalse(compat.forward_compatible(*one_day_after))

    # After exiting context manager, value should be reset.
    self.assertFalse(compat.forward_compatible(*compatibility_date))

 def testGatherNdResourceVariable(self):
   with compat.forward_compatibility_horizon(2019, 4, 30):
     with self.cached_session():
       v = resource_variable_ops.ResourceVariable(
           constant_op.constant([[1, 2], [3, 4], [5, 6]]))
       self.evaluate(variables.global_variables_initializer())
       gather = array_ops.gather_nd(v, [[0, 1], [2, 0]])
       if not context.executing_eagerly():  # .op doesn't make sense in Eager
         self.assertEqual("ResourceGatherNd", gather.op.inputs[0].op.type)
       self.assertAllEqual([2, 5], gather)

  def testStaticRegexFullMatchDelegation(self):
    with compat.forward_compatibility_horizon(2018, 11, 20):
      with self.cached_session():
        input_tensor = constant_op.constant("foo", dtypes.string)
        pattern = "[a-z]*"
        op = string_ops.regex_full_match(input_tensor, pattern)
        self.assertTrue(op.name.startswith("StaticRegexFullMatch"), op.name)

        pattern_tensor = constant_op.constant("[a-z]*", dtypes.string)
        op_vec = string_ops.regex_full_match(input_tensor, pattern_tensor)
        self.assertTrue(op_vec.name.startswith("RegexFullMatch"), op.name)

  def Test(self):
    def CheckGradients(self, a_shape, b_shape):
      self._compare(a_shape, b_shape, dtype, adjoint_a, adjoint_b)

    with compat.forward_compatibility_horizon(2019, 4, 19):
      CheckGradients(self, [1, 5, 2, 3], [7, 1, 3, 2])
      CheckGradients(self, [2, 3], [1, 3, 5])
      CheckGradients(self, [2, 3], [5, 3, 5])
      CheckGradients(self, [5, 2, 5], [5, 3])
      CheckGradients(self, [5, 2, 2, 3], [3, 5])
      CheckGradients(self, [4, 5, 1, 2, 3], [1, 1, 3, 5])
      CheckGradients(self, [1, 2, 1, 4, 2, 1, 3, 4], [3, 2, 1, 1, 1, 2, 4, 2])

  def test_decorator_with_failure(self):
    compatibility_date = self._compatibility_date()
    one_day_after = self._n_days_after(1)

    class DummyError(Exception):
      pass

    try:
      with compat.forward_compatibility_horizon(*one_day_after):
        raise DummyError()
    except DummyError:
      pass  # silence DummyError

    # After exiting context manager, value should be reset.
    self.assertFalse(compat.forward_compatible(*compatibility_date))

  def testNMS3Then2WithScoreThresh(self):
    # Three boxes are selected based on IOU.
    # One is filtered out by score threshold.

    # TODO(b/26783907): The Sort HLO is not implemented on CPU or GPU.
    if self.device in ["XLA_CPU", "XLA_GPU"]:
      return

    with compat.forward_compatibility_horizon(2018, 8, 8):
      boxes_data = [[0, 0, 1, 1], [0, 0.1, 1, 1.1], [0, -0.1, 1, 0.9],
                    [0, 10, 1, 11], [0, 10.1, 1, 11.1], [0, 100, 1, 101]]
      boxes_np = np.array(boxes_data, dtype=np.float32)

      scores_data = [0.9, 0.75, 0.6, 0.95, 0.5, 0.3]
      scores_np = np.array(scores_data, dtype=np.float32)
      max_output_size = 3
      iou_threshold_np = np.array(0.5, dtype=np.float32)
      score_threshold_np = np.array(0.4, dtype=np.float32)

      with self.cached_session() as sess:
        boxes = array_ops.placeholder(boxes_np.dtype, shape=boxes_np.shape)
        scores = array_ops.placeholder(scores_np.dtype, shape=scores_np.shape)
        iou_threshold = array_ops.placeholder(iou_threshold_np.dtype,
                                              iou_threshold_np.shape)
        score_threshold = array_ops.placeholder(score_threshold_np.dtype,
                                                score_threshold_np.shape)
        with self.test_scope():
          selected_indices = image_ops.non_max_suppression_padded(
              boxes=boxes,
              scores=scores,
              max_output_size=max_output_size,
              iou_threshold=iou_threshold,
              score_threshold=score_threshold,
              pad_to_max_output_size=True)
        inputs_feed = {
            boxes: boxes_np,
            scores: scores_np,
            iou_threshold: iou_threshold_np,
            score_threshold: score_threshold_np
        }
        (indices_tf, num_valid) = sess.run(
            selected_indices, feed_dict=inputs_feed)

        self.assertEqual(indices_tf.size, max_output_size)
        self.assertEqual(num_valid, 2)
        self.assertAllClose(indices_tf[:num_valid], [3, 0])

 def testGradGradFloat32(self):
   with compat.forward_compatibility_horizon(2018, 11, 2):
     with self.test_session():
       x = constant_op.constant(
           [-0.9, -0.7, -0.5, -0.3, -0.1, 0.1, 0.3, 0.5, 0.7, 0.9],
           shape=[2, 5],
           name="x")
       y = nn_ops.leaky_relu(x, alpha=0.1, name="leaky_relu")
       z = gradients_impl.gradients(y, x)
       x_init = np.asarray(
           [[-0.9, -0.7, -0.5, -0.3, -0.1], [0.1, 0.3, 0.5, 0.7, 0.9]],
           dtype=np.float32,
           order="F")
       err = gradient_checker.compute_gradient_error(
           x, [2, 5], z[0], [2, 5], x_init_value=x_init)
     print("leaky_relu (float32) gradient of gradient err = ", err)
     self.assertLess(err, 1e-4)

  def testGradGradFloat64(self):
    with compat.forward_compatibility_horizon(2018, 11, 2):
      with self.cached_session():

        def f(x):
          assert x.dtype == dtypes.float64
          with backprop.GradientTape() as tape:
            tape.watch(x)
            y = nn_ops.leaky_relu(x)
          return tape.gradient(y, x)

        x = np.asarray(
            [[-0.9, -0.7, -0.5, -0.3, -0.1], [0.1, 0.3, 0.5, 0.7, 0.9]],
            dtype=np.float64,
            order="F")
        err = gradient_checker_v2.max_error(
            *gradient_checker_v2.compute_gradient(f, [x]))
      print("leaky_relu (float64) gradient of gradient err = ", err)
      self.assertLess(err, 1e-10)

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

    with compat.forward_compatibility_horizon(2018, 8, 4):
      dataset = dataset_ops.Dataset.range(10)
      multi_device_iterator = prefetching_ops.MultiDeviceIterator(
          dataset, ["/cpu:1", "/gpu:0"])
      elem_on_1, elem_on_2 = multi_device_iterator.get_next()

      config = config_pb2.ConfigProto(device_count={"CPU": 2, "GPU": 1})
      with self.test_session(config=config) as sess:
        sess.run(multi_device_iterator.initializer)
        for i in range(0, 10, 2):
          self.assertEqual(i, sess.run(elem_on_1))
          self.assertEqual(i + 1, sess.run(elem_on_2))
        with self.assertRaises(errors.OutOfRangeError):
          sess.run(elem_on_1)
          sess.run(elem_on_2)

  def testNMS3From6Boxes(self):
    with compat.forward_compatibility_horizon(2018, 8, 8):
      # Three boxes are selected based on IOU.
      boxes_data = [[0, 0, 1, 1], [0, 0.1, 1, 1.1], [0, -0.1, 1, 0.9],
                    [0, 10, 1, 11], [0, 10.1, 1, 11.1], [0, 100, 1, 101]]
      boxes_np = np.array(boxes_data, dtype=np.float32)

      scores_data = [0.9, 0.75, 0.6, 0.95, 0.5, 0.3]
      scores_np = np.array(scores_data, dtype=np.float32)

      max_output_size = 3
      iou_threshold_np = np.array(0.5, dtype=np.float32)
      score_threshold_np = np.array(0.0, dtype=np.float32)

      with self.cached_session() as sess:
        boxes = array_ops.placeholder(boxes_np.dtype, shape=boxes_np.shape)
        scores = array_ops.placeholder(scores_np.dtype, shape=scores_np.shape)
        iou_threshold = array_ops.placeholder(iou_threshold_np.dtype,
                                              iou_threshold_np.shape)
        score_threshold = array_ops.placeholder(score_threshold_np.dtype,
                                                score_threshold_np.shape)
        with self.test_scope():
          selected_indices = image_ops.non_max_suppression_padded(
              boxes=boxes,
              scores=scores,
              max_output_size=max_output_size,
              iou_threshold=iou_threshold,
              score_threshold=score_threshold,
              pad_to_max_output_size=True)
        inputs_feed = {
            boxes: boxes_np,
            scores: scores_np,
            score_threshold: score_threshold_np,
            iou_threshold: iou_threshold_np
        }
        (indices_tf, num_valid) = sess.run(
            selected_indices, feed_dict=inputs_feed)

        self.assertEqual(indices_tf.size, max_output_size)
        self.assertEqual(num_valid, 3)
        self.assertAllClose(indices_tf[:num_valid], [3, 0, 5])

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

    with compat.forward_compatibility_horizon(2018, 8, 4):
      host_dataset = dataset_ops.Dataset.range(10)
      device_dataset = host_dataset.apply(
          prefetching_ops.copy_to_device("/gpu:0", source_device="/cpu:0"))
      back_to_cpu_dataset = device_dataset.apply(
          prefetching_ops.copy_to_device("/cpu:0", source_device="/gpu:0"))

      with ops.device("/cpu:0"):
        iterator = back_to_cpu_dataset.make_initializable_iterator()
        next_element = iterator.get_next()

      with self.cached_session() as sess:
        sess.run(iterator.initializer)
        for i in range(10):
          self.assertEqual(i, sess.run(next_element))
        with self.assertRaises(errors.OutOfRangeError):
          sess.run(next_element)

 def test1DTensorAsInputNoReshape(self):
   with compat.forward_compatibility_horizon(2018, 8, 27):
     self._testSoftmax(
         np.array([3., 2., 3., 9.]).astype(np.float64), use_gpu=False)
     self._testOverflow(use_gpu=False)

  def testIteratorStringHandleFuture(self):
    with forward_compat.forward_compatibility_horizon(2018, 8, 4):
      dataset_3 = dataset_ops.Dataset.from_tensor_slices([1, 2, 3])
      dataset_4 = dataset_ops.Dataset.from_tensor_slices([10, 20, 30, 40])

      iterator_3 = dataset_ops.make_one_shot_iterator(dataset_3)
      iterator_4 = dataset_ops.make_one_shot_iterator(dataset_4)

      handle_placeholder = array_ops.placeholder(dtypes.string, shape=[])
      feedable_iterator = iterator_ops.Iterator.from_string_handle(
          handle_placeholder, dataset_ops.get_legacy_output_types(dataset_3),
          dataset_ops.get_legacy_output_shapes(dataset_3))
      next_element = feedable_iterator.get_next()

      self.assertTrue(dataset_ops.get_structure(dataset_3).is_compatible_with(
          dataset_ops.get_structure(feedable_iterator)))
      self.assertTrue(dataset_ops.get_structure(dataset_4).is_compatible_with(
          dataset_ops.get_structure(feedable_iterator)))

      with self.cached_session() as sess:
        iterator_3_handle = sess.run(iterator_3.string_handle())
        iterator_4_handle = sess.run(iterator_4.string_handle())

        self.assertEqual(
            10,
            sess.run(
                next_element,
                feed_dict={handle_placeholder: iterator_4_handle}))
        self.assertEqual(
            1,
            sess.run(
                next_element,
                feed_dict={handle_placeholder: iterator_3_handle}))
        self.assertEqual(
            20,
            sess.run(
                next_element,
                feed_dict={handle_placeholder: iterator_4_handle}))
        self.assertEqual(
            2,
            sess.run(
                next_element,
                feed_dict={handle_placeholder: iterator_3_handle}))
        self.assertEqual(
            30,
            sess.run(
                next_element,
                feed_dict={handle_placeholder: iterator_4_handle}))
        self.assertEqual(
            3,
            sess.run(
                next_element,
                feed_dict={handle_placeholder: iterator_3_handle}))
        self.assertEqual(
            40,
            sess.run(
                next_element,
                feed_dict={handle_placeholder: iterator_4_handle}))
        with self.assertRaises(errors.OutOfRangeError):
          sess.run(
              next_element, feed_dict={handle_placeholder: iterator_3_handle})
        with self.assertRaises(errors.OutOfRangeError):
          sess.run(
              next_element, feed_dict={handle_placeholder: iterator_4_handle})