def testRemoveNonexistentOverrideValue(self):
    stepper = NodeStepper(self.sess, self.e)
    self.assertEqual([], stepper.override_names())

    with self.assertRaisesRegexp(
        ValueError, "No overriding value exists for tensor \"c:0\""):
      stepper.remove_override("c:0")

  def testIsFeedableShouldGiveCorrectAnswers(self):
    stepper = NodeStepper(self.sess, self.e)

    self.assertTrue(stepper.is_feedable("a/read:0"))
    self.assertTrue(stepper.is_feedable("b/read:0"))
    self.assertTrue(stepper.is_feedable("c:0"))
    self.assertTrue(stepper.is_feedable("d:0"))

  def testIsFeedable(self):
    stepper = NodeStepper(self.sess, self.e)

    self.assertTrue(stepper.is_feedable("a/read:0"))
    self.assertTrue(stepper.is_feedable("b/read:0"))
    self.assertTrue(stepper.is_feedable("c:0"))
    self.assertTrue(stepper.is_feedable("d:0"))

  def testUsingNodesNotUsingIntermediateTensors(self):
    stepper = NodeStepper(self.sess, self.e)

    # There should be no handles before any cont() calls.
    self.assertEqual([], stepper.handle_names())

    # Before the cont() call, the stepper should not have access to the value
    # of c:0.
    with self.assertRaisesRegexp(
        ValueError,
        "This stepper instance does not have access to the value of tensor "
        "\"c:0\""):
      stepper.get_tensor_value("c:0")

    # Using the node/tensor itself, instead of the name str, should work on
    # cont().
    result = stepper.cont(self.c)
    self.assertAllClose(6.0, result)
    self.assertEqual({}, stepper.last_feed_types())

    self.assertEqual(["c:0"], stepper.handle_names())

    # After the cont() call, the stepper should have access to the value of c:0
    # via a tensor handle.
    self.assertAllClose(6.0, stepper.get_tensor_value("c:0"))

    result = stepper.cont(self.e)
    self.assertAllClose(24.0, result)
    self.assertEqual({
        "c:0": NodeStepper.FEED_TYPE_HANDLE
    }, stepper.last_feed_types())

  def testTransitiveClosureWithCrossLinksShouldHaveCorrectOrder(self):
    stepper = NodeStepper(self.sess, "z:0")

    sorted_nodes = stepper.sorted_nodes()
    self.assertEqual(4, len(sorted_nodes))
    self.assertLess(sorted_nodes.index("x"), sorted_nodes.index("x/read"))
    self.assertLess(sorted_nodes.index("x"), sorted_nodes.index("y"))
    self.assertLess(sorted_nodes.index("x"), sorted_nodes.index("z"))
    self.assertLess(sorted_nodes.index("y"), sorted_nodes.index("z"))

  def testContWithPlaceholders(self):
    stepper = NodeStepper(
        self.sess,
        self.y,
        feed_dict={
            self.ph0: [[1.0, 2.0], [-3.0, 5.0]],
            self.ph1: [[-1.0], [0.5]]
        })

    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 testContToUpdateA(self):
    stepper = NodeStepper(self.sess, "optim")

    result = stepper.cont("a:0")
    self.assertAllClose(1.0, result)
    self.assertEqual({}, stepper.last_feed_types())

    result = stepper.cont("optim/learning_rate:0")
    self.assertAllClose(0.01, result)
    self.assertEqual({}, stepper.last_feed_types())

    # Before any cont calls on ApplyGradientDescent, there should be no "dirty"
    # variables.
    self.assertEqual(set(), stepper.dirty_variables())

    # First, all the two control inputs to optim.
    result = stepper.cont("optim/update_a/ApplyGradientDescent")

    # Now variable a should have been marked as dirty due to the update
    # by optim/update_a/ApplyGradientDescent.
    self.assertEqual({"a:0"}, stepper.dirty_variables())
    self.assertIsNone(result)
    self.assertEqual({
        "optim/learning_rate:0": NodeStepper.FEED_TYPE_HANDLE
    }, stepper.last_feed_types())

    # Check that Variable "a" has been updated properly, but "b", "c" and "d"
    # remain the same.
    # For backprop on Variable a:
    #   Because f = a * b * b * c, df / da = b * b * c.
    #   1.0 - learning_rate * b * b * c
    #     = 1.0 -  0.01 * 2.0 * 2.0 * 4.0 = 0.84.
    self.assertAllClose(0.84, self.sess.run(self.a))
    self.assertAllClose(2.0, self.sess.run(self.b))
    self.assertAllClose(4.0, self.sess.run(self.c))

  def testRestoreVariableValues(self):
    """Test restore_variable_values() restores the old values of variables."""

    stepper = NodeStepper(self.sess, "optim")

    stepper.cont(
        "optim/update_b/ApplyGradientDescent", restore_variable_values=True)
    self.assertAllClose(1.84, self.sess.run(self.b))

    stepper.restore_variable_values()
    self.assertAllClose(2.0, self.sess.run(self.b))

  def testAttemptToContToPlaceholder(self):
    stepper = NodeStepper(
        self.sess,
        self.y,
        feed_dict={
            self.ph0: [[1.0, 2.0], [-3.0, 5.0]],
            self.ph1: [[-1.0], [0.5]]
        })

    with self.assertRaisesRegexp(ValueError,
                                 r"Should not call cont\(\) on a Placeholder"):
      stepper.cont(self.ph0)

  def testIsPlaceholdersShouldGiveCorrectAnswers(self):
    stepper = NodeStepper(self.sess, self.y)

    self.assertTrue(stepper.is_placeholder(self.ph0.name))
    self.assertTrue(stepper.is_placeholder(self.ph1.name))

    self.assertFalse(stepper.is_placeholder(self.x.name))
    self.assertFalse(stepper.is_placeholder(self.y.name))

    with self.assertRaisesRegexp(ValueError,
                                 "A is not in the transitive closure"):
      self.assertFalse(stepper.is_placeholder("A"))

  def testAttemptToContToFetchNotInTransitiveClosure(self):
    stepper = NodeStepper(self.sess, "e:0")

    self.assertEqual(
        ["a:0", "b:0", "b/read:0", "a/read:0", "c:0", "d:0", "e:0"],
        stepper.sorted_transitive_closure())

    with self.assertRaisesRegexp(
        ValueError,
        "Target \"f:0\" is not in the transitive closure for the fetch of the "
        "stepper: \"e:0\""):
      stepper.cont("f:0")

  def testAttemptToContToPlaceholderWithTensorNameFeedKeysShouldWork(self):

    ph0_feed = [[1.0, 2.0], [-3.0, 5.0]]
    ph1_feed = [[-1.0], [0.5]]
    stepper = NodeStepper(
        self.sess,
        self.y,
        feed_dict={
            self.ph0.name: ph0_feed,
            self.ph1.name: ph1_feed,
        })

    self.assertAllClose(ph0_feed, stepper.cont(self.ph0))
    self.assertEqual({
        self.ph0.name: NodeStepper.FEED_TYPE_CLIENT
    }, stepper.last_feed_types())

    self.assertAllClose(ph1_feed, stepper.cont(self.ph1))
    self.assertEqual({
        self.ph1.name: NodeStepper.FEED_TYPE_CLIENT
    }, stepper.last_feed_types())

    ph0_node = self.sess.graph.as_graph_element("ph0")
    self.assertAllClose(ph0_feed, stepper.cont(ph0_node))
    self.assertEqual({
        self.ph0.name: NodeStepper.FEED_TYPE_CLIENT
    }, stepper.last_feed_types())

    self.assertAllClose([[-1.0], [6.0]], stepper.finalize())

  def testContToNodeWithOutputTensors(self):
    """cont() to an op should cache its output tensors if appropriate."""

    stepper = NodeStepper(self.sess, "optim")

    # In the transitive closure of the stepper, look for an op of which the
    # output tensor also is in the transitive closure.
    # Do not assume a specific op, e.g., ""gradients/e_grad/Reshape_1",
    # because it may vary between builds.
    closure = stepper.sorted_transitive_closure()
    op_with_output_in_closure = None
    for element_name in closure:
      if element_name + ":0" in closure:
        op_with_output_in_closure = str(element_name)
        break

    self.assertIsNotNone(op_with_output_in_closure)
    output_tensor = op_with_output_in_closure + ":0"

    # The op "gradients/?_grad/Reshape_1" is in the transitive closure of the
    # stepper, because it is the control input to another o. However, its
    # output tensor "gradients/?_grad/Reshape_1:0" is also in the transitive
    # closure, because it is the (non-control) input of certain ops. Calling
    # cont() on the op should lead to the caching of the tensor handle for
    # the output tensor.
    stepper.cont(op_with_output_in_closure)

    self.assertEqual([output_tensor], stepper.handle_names())

    # Do a cont() call that uses the cached tensor of
    # "gradients/?_grad/Reshape_1:0".
    stepper.cont(output_tensor)
    self.assertEqual({
        output_tensor: NodeStepper.FEED_TYPE_HANDLE
    }, stepper.last_feed_types())

  def testContToUpdateB(self):
    stepper = NodeStepper(self.sess, "optim")

    result = stepper.cont("optim/update_b/ApplyGradientDescent")
    self.assertIsNone(result)
    self.assertEqual(set(["b:0"]), stepper.dirty_variables())

    # For backprop on Variable b:
    #   Because f = a * b * b * c, df / da = 2 * a * b * c.
    #   2.0 - learning_rate * 2 * a * b * c
    #     = 2.0 - 0.01 * 2 * 1.0 * 2.0 * 4.0 = 1.84
    self.assertAllClose(1.0, self.sess.run(self.a))
    self.assertAllClose(1.84, self.sess.run(self.b))
    self.assertAllClose(4.0, self.sess.run(self.c))

  def testUsingNamesNotUsingIntermediateTensors(self):
    stepper = NodeStepper(self.sess, "e:0")

    # The first cont() call should have used no feeds.
    result = stepper.cont("c:0")
    self.assertAllClose(6.0, result)
    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.assertEqual({
        "c:0": NodeStepper.FEED_TYPE_HANDLE
    }, stepper.last_feed_types())

  def testGetTensorValueWorksOnPlaceholder(self):
    stepper = NodeStepper(
        self.sess,
        self.y,
        feed_dict={
            self.ph0: [[1.0, 2.0], [-3.0, 5.0]],
            self.ph1: [[-1.0], [0.5]]
        })

    self.assertAllClose([[1.0, 2.0], [-3.0, 5.0]],
                        stepper.get_tensor_value("ph0"))
    self.assertAllClose([[1.0, 2.0], [-3.0, 5.0]],
                        stepper.get_tensor_value("ph0:0"))
    with self.assertRaisesRegexp(
        KeyError, r"The name 'ph0:1' refers to a Tensor which does not exist"):
      stepper.get_tensor_value("ph0:1")

  def testFinalize(self):
    """Test finalize() to restore variables and run the original fetch."""

    stepper = NodeStepper(self.sess, "optim")

    # Invoke update_b before calling finalize.
    stepper.cont("optim/update_b/ApplyGradientDescent",
                 restore_variable_values=True)

    result = stepper.finalize()
    self.assertIsNone(result)

    # The results of the Variable updates should be the same as if no cont()
    # call has occurred on update_b.
    self.assertAllClose(0.84, self.sess.run(self.a))
    self.assertAllClose(1.84, self.sess.run(self.b))
    self.assertAllClose(3.96, self.sess.run(self.c))

  def testUpdateTwiceRestoreVariable(self):
    stepper = NodeStepper(self.sess, "optim")

    result = stepper.cont("optim/update_a/ApplyGradientDescent",
                          restore_variable_values=True)
    self.assertIsNone(result)
    self.assertEqual({"a:0"}, stepper.dirty_variables())

    result = stepper.cont("optim/update_b/ApplyGradientDescent",
                          restore_variable_values=True)
    self.assertIsNone(result)
    # Variables a and c should have been restored and hence no longer dirty.
    # Variable b should have been marked as dirty.
    self.assertEqual({"b:0"}, stepper.dirty_variables())

    # The result of the update should be identitcal to as if only update_b is
    # run.
    self.assertAllClose(1.0, self.sess.run(self.a))
    self.assertAllClose(1.84, self.sess.run(self.b))
    self.assertAllClose(4.0, self.sess.run(self.c))

  def testContToFetchNotInTransitiveClosureShouldError(self):
    stepper = NodeStepper(self.sess, "e:0")

    sorted_nodes = stepper.sorted_nodes()
    self.assertEqual(7, len(sorted_nodes))
    self.assertLess(sorted_nodes.index("a"), sorted_nodes.index("a/read"))
    self.assertLess(sorted_nodes.index("b"), sorted_nodes.index("b/read"))
    self.assertLess(sorted_nodes.index("a"), sorted_nodes.index("c"))
    self.assertLess(sorted_nodes.index("b"), sorted_nodes.index("c"))
    self.assertLess(sorted_nodes.index("a"), sorted_nodes.index("d"))
    self.assertLess(sorted_nodes.index("d"), sorted_nodes.index("e"))
    self.assertLess(sorted_nodes.index("c"), sorted_nodes.index("e"))

    self.assertSetEqual(
        {"e:0", "d:0", "c:0", "a/read:0", "b/read:0", "b:0", "a:0"},
        set(stepper.closure_elements()))

    with self.assertRaisesRegexp(
        ValueError,
        "Target \"f:0\" is not in the transitive closure for the fetch of the "
        "stepper"):
      stepper.cont("f:0")

  def testContAfterUpdateWithoutRestoringVariableValue(self):
    stepper = NodeStepper(self.sess, "optim")

    # First, update Variable a from 1.0 to 0.84.
    result = stepper.cont("optim/update_a/ApplyGradientDescent",
                          restore_variable_values=True)
    self.assertIsNone(result)
    self.assertEqual(set(["a:0"]), stepper.dirty_variables())
    self.assertAllClose(0.84, self.sess.run(self.a))
    self.assertAllClose(2.0, self.sess.run(self.b))
    self.assertAllClose(4.0, self.sess.run(self.c))

    # Second, update Variable b without the default restore_variable_values.
    result = stepper.cont(
        "optim/update_b/ApplyGradientDescent", restore_variable_values=False)
    self.assertIsNone(result)
    # For the backprop on Variable b under the updated value of a:
    #   2.0 - learning_rate * 2 * a' * b * c
    #     = 2.0 - 0.01 * 2 * 0.84 * 2.0 * 4.0 = 1.8656
    self.assertAllClose(0.84, self.sess.run(self.a))
    self.assertAllClose(1.8656, self.sess.run(self.b))
    self.assertAllClose(4.0, self.sess.run(self.c))