Function WhileLoopHelper

tensorflow/compiler/xla/client/lib/loops.cc:25–80

namespace xla {

StatusOr<std::vector<XlaOp>> WhileLoopHelper(
    const WhileLoopHelperConditionFunction& condition_function,
    const WhileLoopHelperBodyFunction& body_function,
    absl::Span<const XlaOp> initial_values, absl::string_view name,
    XlaBuilder* builder) {
  int arity = initial_values.size();
  std::vector<Shape> var_shapes;
  var_shapes.reserve(arity);
  for (const XlaOp& input : initial_values) {
    TF_ASSIGN_OR_RETURN(auto shape, builder->GetShape(input));
    var_shapes.push_back(std::move(shape));
  }
  Shape tuple_shape = ShapeUtil::MakeTupleShape(var_shapes);

  // Unpacks a tuple into its component parts.
  auto unpack_tuple = [](XlaOp tuple, int arity, XlaBuilder* builder) {
    std::vector<XlaOp> elements(arity);
    for (int i = 0; i < arity; ++i) {
      elements[i] = GetTupleElement(tuple, i);
    }
    return elements;
  };

  // Build the condition.
  std::unique_ptr<XlaBuilder> cond_builder =
      builder->CreateSubBuilder(absl::StrCat(name, "_condition"));
  {
    auto parameter = Parameter(cond_builder.get(), 0, tuple_shape, "parameter");

    TF_RETURN_IF_ERROR(
        condition_function(unpack_tuple(parameter, arity, cond_builder.get()),
                           cond_builder.get())
            .status());
  }
  TF_ASSIGN_OR_RETURN(auto cond, cond_builder->Build());

  // Build the body.
  std::unique_ptr<XlaBuilder> body_builder =
      builder->CreateSubBuilder(absl::StrCat(name, "_body"));
  {
    auto parameter = Parameter(body_builder.get(), 0, tuple_shape, "parameter");

    TF_ASSIGN_OR_RETURN(
        auto result,
        body_function(unpack_tuple(parameter, arity, body_builder.get()),
                      body_builder.get()));

    TF_RET_CHECK(result.size() == initial_values.size());
    Tuple(body_builder.get(), result);
  }
  TF_ASSIGN_OR_RETURN(auto body, body_builder->Build());

  auto outputs = While(cond, body, Tuple(builder, initial_values));

  return unpack_tuple(outputs, arity, builder);
}

StatusOr<std::vector<XlaOp>> ForEachIndex(