Fix bug in C++ gradient computation where multi-output nodes may not have been processed.

If a subset of the outputs of a multi-output node were included in the
'inputs' argument to AddSymbolicGradients(), they would not be
populated in 'grad_outputs'. This is because the pending count of the
multi-output node would never go to zero (unless the backprop
traversal happened to visit all the output edges, even though not all
were explicitly requested).

This change fixes the bug by processing any remaining requested edges
after finishing the traversal.

PiperOrigin-RevId: 162653892

2 files changed, 53 insertions, 0 deletions

diff --git a/tensorflow/cc/framework/gradients.cc b/tensorflow/cc/framework/gradients.cc
index d454bc3e4c..cec3ebc0ad 100644
--- a/tensorflow/cc/framework/gradients.cc
+++ b/tensorflow/cc/framework/gradients.cc
@@ -352,6 +352,23 @@ Status SymbolicGradientBuilder::AddGradients() {
           BackpropAlongEdge(dx[dx_index++], {e->src(), e->src_output()}));
     }
   }
+
+  // Check if any input nodes still have pending gradients and have not been
+  // processed yet. This happens if not all outputs of a node are in 'inputs_'.
+  std::unordered_map<Node*, int> requested_grads;
+  for (Output nout : inputs_) {
+    if (pending_[nout.node()->id()] > 0) {
+      DCHECK_GT(nout.node()->num_outputs(), 1);
+      int idx = input_nodes_[nout];
+      DCHECK(((*grad_outputs_)[idx].node() == nullptr));
+      TF_RETURN_IF_ERROR(SumGradients(nout, &(*grad_outputs_)[idx]));
+      ++requested_grads[nout.node()];
+    }
+  }
+  for (auto& p : requested_grads) {
+    int num_requested_inputs = p.first->num_outputs() - pending_[p.first->id()];
+    CHECK_EQ(num_requested_inputs, p.second);
+  }
   return Status::OK();
 }
 
diff --git a/tensorflow/cc/framework/gradients_test.cc b/tensorflow/cc/framework/gradients_test.cc
index 2aad978480..24af7d567b 100644
--- a/tensorflow/cc/framework/gradients_test.cc
+++ b/tensorflow/cc/framework/gradients_test.cc
@@ -260,6 +260,42 @@ TEST_F(GradientsTest, StackUnstack_StopBackprop) {
   CompareTestAndExpectedGraphs();
 }
 
+TEST_F(GradientsTest, StackUnstack_SubsetOfUnstackOutputs) {
+  // Constructs an unstack with three outputs, and takes the gradient with
+  // respect to only two of the outputs. Tests that the output gradients are
+  // computed.
+  for (const bool expected : {false, true}) {
+    const Scope& scope = expected ? scope_expected_ : scope_test_;
+    // Construct forward graph.
+    auto c = Const(scope, 1, {3, 4, 2});
+    auto unpack = Unstack(scope, c, 3);
+    auto x = Identity(scope, unpack.output[0]);
+    auto y = Identity(scope, unpack.output[1]);
+    auto z = Identity(scope, unpack.output[2]);
+    TF_ASSERT_OK(scope.status());
+
+    // Construct grad inputs.
+    auto dy = Const(scope, 4, {4, 2});
+    auto dz = Const(scope, 5, {4, 2});
+
+    if (expected) {
+      // Construct backward graph.
+      auto g1 = Identity(scope, dy);
+      auto g2 = Identity(scope, dz);
+    } else {
+      // Call AddSymbolicGradients.
+      std::vector<Output> grad_outputs;
+      TF_ASSERT_OK(AddSymbolicGradients(scope, {y, z},
+                                        {unpack.output[1], unpack.output[2]},
+                                        {dy, dz}, &grad_outputs));
+      ASSERT_EQ(grad_outputs.size(), 2);
+      EXPECT_TRUE(grad_outputs[0].node() != nullptr);
+      EXPECT_TRUE(grad_outputs[1].node() != nullptr);
+    }
+  }
+  CompareTestAndExpectedGraphs();
+}
+
 TEST_F(GradientsTest, DependentGradOutputs) {
   // Tests that dependent gradients (in this case the gradients w.r.t to the
   // output and one input of MatMul) are computed properly.