Function AddWhileGradientLoop

tensorflow/cc/framework/while_gradients.cc:136–175

Creates the main backprop loop that computes the gradient of the loop associated with `while_ctx`. `grad_inputs` are the partial derivatives w.r.t. the loop outputs, i.e. the exit nodes. `backprop_execution_pred` is the predicate to use for the backprop loop (see AddBackPropLoopCounter()). The partial derivatives w.r.t. the loop inputs, i.e. the input loop vars, are returned in `grad_outputs`.

// The partial derivatives w.r.t. the loop inputs, i.e. the input loop vars, are
// returned in `grad_outputs`.
Status AddWhileGradientLoop(WhileContext* while_ctx,
                            const std::vector<Output>& grad_inputs,
                            const Output& backprop_execution_pred,
                            const Scope& parent_scope,
                            std::vector<Output>* grad_outputs) {
  DCHECK_EQ(grad_inputs.size(), while_ctx->body_outputs().size());
  DCHECK_EQ(while_ctx->body_inputs().size(), while_ctx->body_outputs().size());

  Scope scope = parent_scope.NewSubScope("while");

  // Create while loop:
  //   while backprop_execution_pred:
  //     forward loop body gradient

  // Condition function that returns 'backprop_execution_pred'.
  CondGraphBuilderFn cond_fn = [backprop_execution_pred](
                                   const Scope& scope,
                                   const std::vector<Output>& inputs,
                                   Output* output) {
    *output = backprop_execution_pred;
    return Status::OK();
  };

  // Body function that builds while body gradient subgraph.
  BodyGraphBuilderFn body_fn = [while_ctx](const Scope& scope,
                                           const std::vector<Output>& inputs,
                                           std::vector<Output>* outputs) {
    std::vector<Output> body_outputs =
        ToOutputVector(while_ctx->body_outputs());
    std::vector<Output> body_inputs = ToOutputVector(while_ctx->body_inputs());
    return AddSymbolicGradients(scope, body_outputs, body_inputs, inputs,
                                outputs);
  };

  string frame_name = BackPropFrameName(while_ctx->frame_name());
  TF_RETURN_IF_ERROR(BuildWhileLoop(scope, grad_inputs, cond_fn, body_fn,
                                    frame_name, grad_outputs,
                                    /* create_while_ctx */ false));
  return Status::OK();
}

}  // namespace