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,

                                           new_result_types, new_val_operands);

  // Move all regions from the old `mhlo.while` op to its replacement.
  new_while.getCond().takeBody(region_while.getCond());
  new_while.getBody().takeBody(region_while.getBody());

  // Forward result from old `mhlo.while` with replacement.
  SmallVector<Value> old_while_results = region_while.getResults();

// condition of whether to continue to next iteration.
void PopulateDatasetWhileCond(OpBuilder builder, WhileRegionOp dataset_while,
                              Location loc) {
  auto& cond_region = dataset_while.getCond();
  Block* cond_block = builder.createBlock(&cond_region);
  auto while_input_types = dataset_while.getOperandTypes();
  cond_block->addArguments(

  Operation* op_inst = op.getOperation();
  Location loc = op_inst->getLoc();

  OpBuilder builder(op_inst);

  // Lower the condition to a boolean value (i1).
  Value cond_i1 = LowerCondition(loc, op.getCond(), &builder);
  if (!cond_i1) return failure();

  // Split the basic block before the 'if'.  The new dest will be our merge
  // point.
  Block* orig_block = op_inst->getBlock();

    return failure();
  }
  auto output_buffer_to_size = AddTensorListSizesToTerminator<TF::YieldOp>(
      body_region.front(), *buffer_to_size);

  // Rewrite cond.
  Region& cond_region = while_op.getCond();
  modify_region_arguments(cond_region);
  if (failed(DecomposeTensorListOpsInternal(
          &cond_region.front(), module, buffer_to_size,
          decomposed_partitioned_call_callees))) {
    return failure();

    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;
  auto mirrored_variable_indices_attr =
      replicate->getAttrOfType<ArrayAttr>(kMirroredVariableIndicesAttr);

        }

        if (auto while_op = llvm::dyn_cast<mlir::TF::WhileRegionOp>(owner)) {
          const int operand_number = use.getOperandNumber();
          next_values_to_visit.push_back(
              while_op.getCond().front().getArgument(operand_number));
          next_values_to_visit.push_back(
              while_op.getBody().front().getArgument(operand_number));
          continue;
        }