From 574bcdab42c8db34f0f082ffba69aacbea36c36d Mon Sep 17 00:00:00 2001
From: Yu Yang <yuyang18@baidu.com>
Date: Fri, 23 Feb 2018 13:13:38 +0800
Subject: [PATCH] Add comments

---
 paddle/fluid/operators/while_op.cc | 22 ++++++++++++----------
 1 file changed, 12 insertions(+), 10 deletions(-)

diff --git a/paddle/fluid/operators/while_op.cc b/paddle/fluid/operators/while_op.cc
index 5f51a273dd..8b62b242cf 100644
--- a/paddle/fluid/operators/while_op.cc
+++ b/paddle/fluid/operators/while_op.cc
@@ -261,35 +261,37 @@ class WhileGradOpDescMaker : public framework::SingleGradOpDescMaker {
     for (auto &o : Output(kOutputs)) {
       block_ins.insert(o);
     }
-    std::unordered_set<std::string> extra_inputs;
+    std::unordered_set<std::string> output_grads;
     for (const auto *op : grad_block->AllOps()) {
       for (auto &input_name : op->InputArgumentNames()) {
         // If the input of Op has been recorded or is generated by the forward
         // block, do not make it as input again.
 
+        // The input is located in I/O or other op's outputs or the variable is
+        // located in grad_block's parents
         if (block_ins.find(input_name) != block_ins.end() ||
-            fwd_block->FindVar(input_name) != nullptr ||
-            parent_block->FindVar(input_name) != nullptr) {
+            (fwd_block->FindVarRecursive(input_name) != nullptr ||
+             parent_block->FindVarRecursive(input_name) != nullptr)) {
           continue;
         }
-        extra_inputs.insert(input_name);
+        output_grads.insert(input_name);
       }
       for (auto &output_name : op->OutputArgumentNames()) {
         block_ins.insert(output_name);
       }
     }
 
-    std::vector<std::string> extra_inputs_list;
-    extra_inputs_list.resize(extra_inputs.size());
-    std::copy(extra_inputs.begin(), extra_inputs.end(),
-              extra_inputs_list.begin());
-    while_grad->SetInput(framework::GradVarName(kOutputs), extra_inputs_list);
+    std::vector<std::string> output_grads_list;
+    output_grads_list.resize(output_grads.size());
+    std::copy(output_grads.begin(), output_grads.end(),
+              output_grads_list.begin());
+    while_grad->SetInput(framework::GradVarName(kOutputs), output_grads_list);
 
     while_grad->SetAttrMap(this->Attrs());
     while_grad->SetBlockAttr(kStepBlock, *grad_block);
     // record the original output gradient names, since the gradient name of
     // while operator could be renamed.
-    while_grad->SetAttr("original_output_grad", extra_inputs_list);
+    while_grad->SetAttr("original_output_grad", output_grads_list);
 
     return std::unique_ptr<framework::OpDesc>(while_grad);
   }
-- 
GitLab