Function BuildWrappedBody

tensorflow/compiler/tf2xla/kernels/while_op.cc:213–244

}

xla::StatusOr<xla::XlaComputation> BuildWrappedBody(
    XlaOpKernelContext* ctx, const XlaCompiler::CompilationResult& body,
    const std::vector<bool>& compile_time_const_arg_indices,
    int num_compile_time_const_args, bool has_token_input_output) {
  if (num_compile_time_const_args <= 0) {
    return xla::XlaComputation(body.computation->proto());
  }
  xla::XlaComputation body_wrapper;
  std::unique_ptr<xla::XlaBuilder> cb =
      ctx->builder()->CreateSubBuilder("body_wrapper");
  xla::Shape body_input_shape = body.xla_input_shapes[0];
  auto inputs = xla::Parameter(cb.get(), 0, body_input_shape, "inputs");
  // Call the original body function which has mismatched inputs and outputs
  // and strip the compile time consts from the list of outputs. While requires
  // the inputs and outputs of its body function to match.
  auto outputs = xla::Call(cb.get(), *body.computation, {inputs});
  std::vector<xla::XlaOp> non_compile_time_const_outputs;
  for (int i = 0; i < compile_time_const_arg_indices.size(); i++) {
    if (!compile_time_const_arg_indices[i]) {
      non_compile_time_const_outputs.push_back(
          xla::GetTupleElement(outputs, i));
    }
  }
  // If `body` has a token output, append it to
  // `non_compile_time_const_outputs`.
  if (has_token_input_output) {
    non_compile_time_const_outputs.push_back(
        xla::GetTupleElement(outputs, ctx->num_outputs()));
  }
  xla::Tuple(cb.get(), non_compile_time_const_outputs);
  return cb->Build();
}

xla::XlaOp BuildWhile(XlaOpKernelContext* ctx,
                      const xla::XlaComputation& wrapped_cond,