Method ProcessWhileLoop

tensorflow/cc/framework/gradients.cc:390–435

}

Status SymbolicGradientBuilder::ProcessWhileLoop(Node* exit_node,
                                                 const Output& summed_grads) {
  // TODO(skyewm): detect second-order gradient and return bad status
  // TODO(skyewm): handle (or at least detect) nested while loops

  // TODO(skyewm): handle NoGradient in while loop
  if (summed_grads == NoGradient()) {
    return errors::Unimplemented(
        "Missing gradient into while loop not yet implemented");
  }

  DCHECK(exit_node->IsExit());
  WhileContext* while_ctx = exit_node->while_ctx();
  DCHECK(while_ctx != nullptr);

  // Record 'summed_grads' as the backprop input associated with 'exit_node'
  std::map<Node*, Output>& backprops = while_backprops_[while_ctx];
  DCHECK(backprops.find(exit_node) == backprops.end());
  backprops[exit_node] = summed_grads;

  // Wait until we have all exit nodes' backprops collected before processing
  // the while loop.
  // TODO(skyewm): what if not all the exit nodes are reachable?
  if (backprops.size() < while_ctx->exit_nodes().size()) return Status::OK();

  // We've seen all the exit nodes for this loop and have collected all the
  // backprops. Create the gradient graph for the while loop.
  Scope while_scope =
      scope_.NewSubScope(strings::StrCat(while_ctx->frame_name(), "_grad"));
  std::vector<Output> dy;
  for (Node* n : while_ctx->exit_nodes()) dy.push_back(backprops[n]);
  std::vector<Output> dx;
  TF_RETURN_IF_ERROR(AddWhileLoopGradient(while_ctx, while_scope, dy, &dx));

  // Backprop along the in edges to the while loop (i.e. the inputs to the enter
  // nodes)
  DCHECK_EQ(dx.size(), while_ctx->enter_nodes().size());
  for (int i = 0; i < dx.size(); ++i) {
    Node* enter_node = while_ctx->enter_nodes()[i];
    for (const Edge* e : enter_node->in_edges()) {
      if (e->IsControlEdge()) continue;
      TF_RETURN_IF_ERROR(BackpropAlongEdge(dx[i], {e->src(), e->src_output()}));
    }
  }
  return Status::OK();
}

Status SymbolicGradientBuilder::AddGradients() {
  // Initialize backprops.