Function AddBackPropLoopCounter

tensorflow/cc/framework/while_gradients.cc:96–128

Creates a loop that executes `loop_count` times. The returned output is the boolean predicate indicating if the loop is still executing. This is used to drive the gradient computation for the while loop associated with `while_ctx`.

// drive the gradient computation for the while loop associated with
// `while_ctx`.
Status AddBackPropLoopCounter(WhileContext* while_ctx, const Output& loop_count,
                              const Scope& scope,
                              Output* backprop_execution_pred) {
  // Create while loop:
  //   n = loop_count
  //   while n > 0:
  //     --n

  // Condition function that returns input > 0.
  CondGraphBuilderFn cond_fn = [](const Scope& scope,
                                  const std::vector<Output>& inputs,
                                  Output* output) {
    DCHECK_EQ(inputs.size(), 1);
    *output = ops::Greater(scope, inputs[0], 0);
    return scope.status();
  };

  // Body function that subtracts one from input.
  BodyGraphBuilderFn body_fn = [](const Scope& scope,
                                  const std::vector<Output>& inputs,
                                  std::vector<Output>* outputs) {
    DCHECK_EQ(inputs.size(), 1);
    outputs->emplace_back(ops::Subtract(scope, inputs[0], 1));
    return scope.status();
  };

  string frame_name = BackPropFrameName(while_ctx->frame_name());
  std::vector<Output> outputs;
  TF_RETURN_IF_ERROR(BuildWhileLoop(
      scope, {loop_count}, cond_fn, body_fn, frame_name, &outputs,
      /* create_while_ctx */ false, backprop_execution_pred));
  return Status::OK();
}

// Creates the main backprop loop that computes the gradient of the loop
// associated with `while_ctx`. `grad_inputs` are the partial derivatives