llvm::SmallVector<std::pair<int64_t, llvm::SmallVector<Value, 4>>, 4>
AnnotateCompileOpAndGetExecuteArgToWhileArgsMapping(
    TF::WhileRegionOp while_op, tf_device::ReplicateOp replicate,
    TF::TPUExecuteAndUpdateVariablesOp execute,
    tf_device::LaunchOp compile_launch) {
  Region& body = while_op.getBody();
  Region& cond = while_op.getCond();

  llvm::SmallVector<std::pair<int64_t, llvm::SmallVector<Value, 4>>, 4> mapping;

  for (int64_t i = 0; i < while_op.getNumOperands(); ++i) {
    if (ta_arg_buffer_type(i)) {
      while_op.getResult(i).replaceAllUsesWith(while_op.getOperand(i));
    } else {
      while_op.getResult(i).replaceAllUsesWith(new_while.getResult(i));
    }
  }
  while_op.erase();
  return success();
}

LogicalResult HandleIfOp(TF::IfOp if_op, ModuleOp module,

  auto loc = NameLoc::get(
      builder.getStringAttr("chain_control_outputs@" + while_op.getBody()));
  IslandOp const_wrapper = GetDummyConstant(builder, const_type, loc);

  // Get new operand and result types.
  auto new_operands = llvm::to_vector<4>(while_op->getOperands());
  auto new_result_types = llvm::to_vector<4>(while_op->getResultTypes());
  Value const_output = const_wrapper.getOutputs()[0];

  NameAttrList cond_fn, body_fn;
  cond_fn.set_name("f3");
  body_fn.set_name("f2");
  auto while_op =
      ops::While(s.WithOpName("while"),
                 std::initializer_list<Input>{arg0, arg1}, cond_fn, body_fn);
  auto ret2 = ops::_Retval(s.WithOpName("ret2"), while_op.output[0], 2);
  auto ret3 = ops::_Retval(s.WithOpName("ret3"), while_op.output[1], 3);
  std::unique_ptr<Graph> g(new Graph(OpRegistry::Global()));

    // Propagation to callees.
    auto walk_res = func_op.walk([&](Operation* op) {
      if (auto while_op = dyn_cast<WhileOp>(op)) {
        if (failed(propagate_operands_to_callee_arguments(
                while_op, while_op.getOperands(),
                {while_op.body_function(), while_op.cond_function()},
                func_res)))
          return WalkResult::interrupt();

        PropagateInputToOutput(op->getOperand(result.getResultNumber()),
                               result);
    } else if (auto while_op = dyn_cast<WhileOp>(op)) {
      AnalyzeWhileLoop(while_op, backtrack_analysis.GetAnalysisForFunc(
                                     while_op.body_function()));
    } else if (auto while_region = dyn_cast<WhileRegionOp>(op)) {

  module.walk([&](mlir::TFL::WhileOp while_op) {
    auto cond = symbol_table.lookup<mlir::func::FuncOp>(
        while_op->getAttr("cond").cast<mlir::FlatSymbolRefAttr>().getValue());
    AddCallOpInWhileOpRegion(while_op.getCond(), cond);
    while_op->removeAttr("cond");
    auto body = symbol_table.lookup<mlir::func::FuncOp>(
        while_op->getAttr("body").cast<mlir::FlatSymbolRefAttr>().getValue());

  // Recreate the while op.
  OpBuilder builder(while_op);
  // Now use the filtered original operands, which will be replaced by
  // AddLoadsStoresOutsideControlFlowOp().
  auto new_while = builder.create<TF::WhileOp>(
      while_op.getLoc(), body.getFunctionType().getResults(),
      FilterRange<Value, OperandRange>(while_op.getOperands(),
                                       resource_arg_uses),

      MarkOutsideCompiled(host_if.getOperation());
    }
    if (auto while_op = llvm::dyn_cast<mlir::TF::WhileRegionOp>(op)) {
      if (!HasOutsideCompilationNested(op)) return WalkResult::advance();
      OpBuilder builder(while_op);
      auto host_while = CloneEmptyWhile(while_op.getParallelIterations(),
                                        while_op.getLoc(), builder);
      const auto condition_send_recv_key =

                                TF::WhileOp* while_op) {
  // Given a while loop body function 'body_func', find the tf.While Op that
  // uses it.
  auto uses_optional = body_func.getSymbolUses(module);
  if (!uses_optional.has_value()) {
    body_func.emitOpError() << "no use of while loop body";
    return LogicalResult::failure();
  }
  *while_op = nullptr;
  for (auto& use : uses_optional.value()) {