Function AddForwardLoopCounter

tensorflow/cc/framework/while_gradients.cc:57–90

Creates a loop that counts the number of iterations performed by the while loop associated with `while_ctx`. The returned output yields the iteration count.

// while loop associated with `while_ctx`. The returned output yields the
// iteration count.
Status AddForwardLoopCounter(WhileContext* while_ctx, const Scope& scope,
                             Output* count) {
  // Create while loop:
  //   i = 0
  //   while forward loop predicate is true:
  //     ++i

  Output zero = ops::Const(scope, 0, {});

  // Condition function that returns condition output from original while loop.
  CondGraphBuilderFn cond_fn = [while_ctx](const Scope& scope,
                                           const std::vector<Output>& inputs,
                                           Output* output) {
    *output = ToOutput(while_ctx->cond_output());
    return Status::OK();
  };

  // Body function that adds one to 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::Add(scope, inputs[0], 1));
    return scope.status();
  };

  // Note that this loop runs in the same execution frame as the forward loop.
  std::vector<Output> outputs;
  TF_RETURN_IF_ERROR(BuildWhileLoop(scope, {zero}, cond_fn, body_fn,
                                    while_ctx->frame_name(), &outputs,
                                    /* create_while_ctx */ false));
  *count = outputs[0];
  return Status::OK();
}

// 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