while_op, while_op.getCondAttr());
      func::FuncOp body = SymbolTable::lookupNearestSymbolFrom<func::FuncOp>(
          while_op, while_op.getBodyAttr());
      for (auto &arg : while_op->getOpOperands()) {
        BlockArgument cond_arg = cond.getArgument(arg.getOperandNumber());
        this->join(getLatticeElement(cond_arg), *getLatticeElement(arg.get()));

}

LogicalResult ConvertWhileOp(WhileOp while_op, bool allow_passthrough_args) {
  auto while_region = OpBuilder(while_op).create<TF::WhileRegionOp>(
      while_op.getLoc(), while_op.getResultTypes(), while_op.getInput(),
      while_op.getParallelIterations(), while_op.getIsStateless(),
      while_op.getShapeInvariant());
  CopyDeviceAndUnderscoredAttributes(while_op, while_region);

  YieldOp cond_yield =

// RUN: tf-opt -loop-invariant-code-motion %s -o - | FileCheck %s

// CHECK: while_1([[ARG0:%[^ :]*]]: tensor<i32>, [[ARG1:%[^ :]*]]: tensor<1xf32>)
func.func @while_1(%arg0: tensor<i32>, %arg1: tensor<1xf32>) -> tensor<1xf32> {
  // CHECK: [[CST:%[^ ]*]] = arith.constant dense<1> : tensor<i32>
  // CHECK: "tfl.while"([[ARG0]], [[ARG1]])
  // CHECK: (tensor<i32>, tensor<1xf32>) -> (tensor<i32>, tensor<1xf32>)
  %0:2 = "tfl.while"(%arg0, %arg1) (
    // cond
    {

};

// Prunes result defining op if possible, returns true if pruning was done.
bool TryPruneResultDefiningOp(TF::WhileRegionOp while_op, OpResult result) {
  // Don't prune if result is used.
  if (!result.use_empty()) return false;

  Block& body_block = while_op.getBody().front();
  Block& cond_block = while_op.getCond().front();
  Operation* body_yield_op = body_block.getTerminator();

  func.setPrivate();
}

void RunOnWhile(TF::WhileOp while_op) {
  Operation* op = while_op.getOperation();
  // Create new TFL While op that will be used to replace TF While op.
  auto new_op = OpBuilder(op).create<TFL::WhileOp>(
      op->getLoc(), op->getResultTypes(), op->getOperands(),
      while_op.getIsStateless());
  Location loc = while_op->getLoc();
  CreateRegionWithCall(while_op.cond_function(), new_op.getCond(), loc);

      id_map[i] = j++;
      new_operands.push_back(while_op.getOperand(i));
      new_result_types.push_back(while_op.getResultTypes()[i]);
    }
  }

  auto new_while_op = OpBuilder(while_op).create<WhileOp>(
      while_op.getLoc(), new_result_types, new_operands, while_op->getAttrs());
  new_while_op.getCond().takeBody(while_op.getCond());
  new_while_op.getBody().takeBody(while_op.getBody());

  for (auto i = 0; i < n; ++i) {

      read_var->erase();
    } else if (auto while_op = dyn_cast<TF::WhileOp>(user_op)) {
      GetUpdatedWhileOp<TF::WhileOp>(
          while_op, while_op.cond_function().getArgumentTypes(), args_to_erase);
      while_op->erase();
    } else if (auto while_op = dyn_cast<TF::WhileRegionOp>(user_op)) {
      auto new_while_op = GetUpdatedWhileOp(
          while_op, while_op.getCond().getArgumentTypes(), args_to_erase);

  // Create the new while op.
  auto new_while_operands = llvm::to_vector<8>(while_op.getOperands());
  OpBuilder builder(while_op);
  assert(while_op.getNumOperands() == while_op.getNumResults());
  for (int64_t i = 0; i < while_op.getNumResults(); ++i) {
    auto it = data_var_to_size_var.find(while_op.getOperand(i));
    if (it == data_var_to_size_var.end()) continue;

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;

// RUN: tf-opt -split-input-file -loop-invariant-code-motion %s | FileCheck %s

// CHECK: while_1([[ARG0:%[^ :]*]]: tensor<i32>, [[ARG1:%[^ :]*]]: tensor<1xf32>)
func.func @while_1(%arg0: tensor<i32>, %arg1: tensor<1xf32>) -> tensor<1xf32> {
  // CHECK: [[CST:%[^ ]*]] = arith.constant dense<1> : tensor<i32>
  // CHECK: "tf.WhileRegion"([[ARG0]], [[ARG1]])
  // CHECK: (tensor<i32>, tensor<1xf32>) -> (tensor<i32>, tensor<1xf32>)
  %0:2 = "tf.WhileRegion"(%arg0, %arg1) (