auto num_loop_carried = while_op.getCond().getNumArguments();
  auto not_carried_operands =
      while_op.getOperands().drop_front(num_loop_carried);
  extern_values.insert(not_carried_operands.begin(),
                       not_carried_operands.end());
  auto old_extern_values_size = extern_values.size();

  llvm::SmallVector<Region*, 2> regions{&while_op.getCond(),

  Block &cond = while_op.getCond().front();
  Block &body = while_op.getBody().front();

  auto n = while_op.getNumOperands();
  if (!AllOperationSafe(cond) || !AllOperationSafe(body)) return false;

  // Find all Consumed indices.
  // i is consumed element if result(i) is used outside whileOp or
  // argument(i) is used in whileOp.getCond().
  for (auto i = 0; i < n; ++i) {

      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);
      new_while_op.getCond().takeBody(while_op.getCond());
      new_while_op.getBody().takeBody(while_op.getBody());
      llvm::BitVector erase_indices;
      UpdateTerminatorArguments(new_while_op.getBody(), args_to_erase,

  return to_bool;
}

// Transform a functional IfOp to a region based IfRegionOp.
LogicalResult ConvertIfOp(IfOp if_op) {
  Value cond = ConvertConditionToBoolean(if_op, if_op.getCond());
  OpBuilder builder(if_op);
  auto if_region = builder.create<TF::IfRegionOp>(
      if_op.getLoc(), if_op.getResultTypes(), cond, if_op.getIsStateless(),
      builder.getStringAttr(if_op.then_function().getName()),

    regions.push_back(
        RegionSuccessor(&getCond(), getCond().front().getArguments()));
  } else if (point.isParent()) {
    // The parent branches to 'cond'. It is also considered to branch to `body`
    // in case the terminator of `cond` doesn't forward the arguments of `cond`.
    regions.push_back(
        RegionSuccessor(&getCond(), getCond().front().getArguments()));
    regions.push_back(

        worklist, /*extern_values_passthrough=*/false,
        /*only_one_return_value=*/false);
  }

  // Look through ToBool operations for the condition.
  Value cond = if_region.getCond();
  auto to_bool = dyn_cast_or_null<ToBoolOp>(cond.getDefiningOp());
  if (to_bool) cond = to_bool.getOperand();

  // Once we have the `then` and `else` functions ready (either outlined or

      rewriter.eraseOp(op.getOperation());
      return success();
    }

    // Extract the constant cond value.
    DenseElementsAttr cond;
    if (!matchPattern(op.getCond(), m_Constant(&cond))) return failure();

    // TODO(hinsu): Handle constants that are not scalar booleans.
    auto cond_type = mlir::dyn_cast<RankedTensorType>(cond.getType());

                                                while_op.getInput());
      }
      return;
    }
    if (auto while_op = dyn_cast<TF::WhileRegionOp>(op)) {
      PropagatePotentiallyWrittenUpFromCallee(while_op.getCond(),
                                              while_op.getInput());
      PropagatePotentiallyWrittenUpFromCallee(while_op.getBody(),
                                              while_op.getInput());
      return;

// and retuns values corresponding to that result are dropped.
LogicalResult CanonicalizeWhileRegion(TF::WhileRegionOp op) {
  Region &body = op.getBody();
  Region &cond = op.getCond();
  llvm::BitVector can_eliminate(op.getNumResults());

  // Traverse in reverse order so that indices to be deleted stay unchanged.
  for (OpResult result : llvm::reverse(op.getResults())) {

                        communication_key_index++)
              .str();
      auto& cond = host_while.getCond();
      cond.push_back(new Block);
      auto condition =
          while_op.getCond().front().getTerminator()->getOperand(0);
      builder.setInsertionPoint(while_op.getCond().front().getTerminator());
      builder.create<mlir::TF::XlaSendToHostOp>(while_op.getLoc(), condition,