caller->getLoc(), ArrayRef<Type>{},
        ArrayRef<Value>{operand, caller->getResult(updated_index)});
  }
}

// Lifts loads/stores from while loop's body and cond functions.
LogicalResult HandleWhileLoop(TF::WhileOp while_op, func::FuncOp body,
                              func::FuncOp cond) {
  auto return_op = body.front().getTerminator();
  llvm::SmallDenseMap<int64_t, ResourceArgUseInfo> body_use_info;