Function BuildWhileLoop

tensorflow/cc/ops/while_loop.cc:172–247

A while loop with a single loop variable looks like this: (output) ^ +---------------+ | | body subgraph +-------------+ Exit +---------------+ | ^ ^ | | | | Switch<--------+ v ^ | NextIteration | +------+--------+ | +---->| cond subgraph | | | +----------

// duplicated for each loop variable.
// TODO(skyewm): link to public version of design doc
Status BuildWhileLoop(const Scope& scope, const std::vector<Output>& inputs,
                      const CondGraphBuilderFn& cond,
                      const BodyGraphBuilderFn& body, const string& frame_name,
                      OutputList* outputs, bool create_while_ctx,
                      Output* cond_output) {
  DCHECK(!inputs.empty());
  DCHECK(outputs != nullptr);
  DCHECK(outputs->empty());

  TF_RETURN_IF_ERROR(scope.status());
  const size_t num_loop_vars = inputs.size();

  std::vector<Output> enter_outputs(num_loop_vars);
  for (int i = 0; i < num_loop_vars; ++i) {
    enter_outputs[i] = internal::Enter(scope, inputs[i], frame_name);
  }
  TF_RETURN_IF_ERROR(scope.status());

  std::vector<Output> merge_outputs(num_loop_vars);
  for (int i = 0; i < num_loop_vars; ++i) {
    TF_RETURN_IF_ERROR(
        CreateMerge(scope, i, enter_outputs[i], &merge_outputs[i]));
  }

  Output cond_out;
  TF_RETURN_IF_ERROR(CreateCond(scope, cond, merge_outputs, &cond_out));
  if (cond_output != nullptr) *cond_output = cond_out;

  std::vector<Output> switch_trues(num_loop_vars);
  std::vector<Output> switch_falses(num_loop_vars);
  for (int i = 0; i < num_loop_vars; ++i) {
    auto switch_i = Switch(scope, merge_outputs[i], cond_out);
    switch_trues[i] = switch_i.output_true;
    switch_falses[i] = switch_i.output_false;
  }
  TF_RETURN_IF_ERROR(scope.status());

  std::vector<Output> body_outputs;
  TF_RETURN_IF_ERROR(CreateBody(scope, body, switch_trues, &body_outputs));

  std::vector<Output> next_outputs(num_loop_vars);
  for (int i = 0; i < num_loop_vars; ++i) {
    next_outputs[i] = NextIteration(scope, body_outputs[i]);
    DCHECK_EQ(next_outputs[i].node()->name(), NextIterationName(scope, i));
  }
  TF_RETURN_IF_ERROR(scope.status());

  // Create the backedges from the NextIteration nodes to the Merge nodes.
  for (int i = 0; i < num_loop_vars; ++i) {
    const int merge_backedge_output_index = 1;
    scope.graph()->AddEdge(next_outputs[i].node(), next_outputs[i].index(),
                           merge_outputs[i].node(),
                           merge_backedge_output_index);
  }

  outputs->resize(num_loop_vars);
  for (int i = 0; i < num_loop_vars; ++i) {
    (*outputs)[i] = internal::Exit(scope, switch_falses[i]);