// CHECK-NEXT:       type: INT32,
// CHECK-NEXT:       buffer: 3,
// CHECK-NEXT:       name: "WhileOp",
// CHECK-NEXT:       quantization: {
// CHECK-EMPTY:
// CHECK-NEXT:       },
// CHECK-NEXT:       has_rank: true
// CHECK-NEXT:     }, {
// CHECK-NEXT:       shape: [ 1 ],
// CHECK-NEXT:       buffer: 4,
// CHECK-NEXT:       name: "WhileOp1",
// CHECK-NEXT:       quantization: {
// CHECK-EMPTY:
// CHECK-NEXT:       },

        insert(grad.getHandle(), grad.getSource().str(), func_block);
      } else if (auto while_op = llvm::dyn_cast<TF::WhileOp>(op)) {
        for (const auto& entry : AccessedGradients(
                 {while_op.body_function(), while_op.cond_function()}, module))
          for (const string& source : entry.getSecond())
            insert(while_op.getOperand(entry.getFirst()), source, func_block);

          } else if (auto while_op =
                         llvm::dyn_cast_or_null<TF::WhileOp>(next_op)) {
            func::FuncOp func = while_op.body_function();
            auto func_argument = func.getArgument(next_op_operand_num);
            // Check if the op is returned without mutation. Returning values
            // from a while op follow return or identity -> return pattern.

    return success();
  }
  // Create the new while op.
  auto new_while_operands = llvm::to_vector<8>(while_op.getOperands());
  for (int64_t i = 0; i < while_op.getNumResults(); ++i) {
    auto it = buffer_to_size->find(while_op.getOperand(i));
    if (it == buffer_to_size->end()) continue;
    new_while_operands.push_back(it->getSecond().size);
  }
  auto new_while = builder.create<TF::WhileOp>(

    } else if (auto while_op = dyn_cast<TF::WhileRegionOp>(op)) {
      for (auto &region : while_op->getRegions()) {
        for (auto [arg, value] :
             llvm::zip(region.getArguments(), while_op->getOperands())) {
          this->join(getLatticeElement(arg), *getLatticeElement(value));
        }
      }
    } else if (auto while_op = dyn_cast<TF::WhileOp>(op)) {

    // 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();

    while_op->removeAttr("cond");
    auto body = symbol_table.lookup<mlir::func::FuncOp>(
        while_op->getAttr("body").cast<mlir::FlatSymbolRefAttr>().getValue());
    AddCallOpInWhileOpRegion(while_op.getBody(), body);
    while_op->removeAttr("body");
  });
}

void AddRegionsForStableHLOOp(mlir::ModuleOp module) {
  mlir::SymbolTable symbol_table(module);

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;

      IsSingleCallRegion(while_region.getBody()), while_arg_matcher);

  // All existing inputs to while region are inputs to the functional while.
  auto new_inputs = llvm::to_vector<4>(while_region.getOperands());

  // All existing results will also be generated by the functional while.
  auto new_result_types = llvm::to_vector<4>(while_region.getResultTypes());

  std::string cond_name, body_name;

        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)) {