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(

    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;
        }

  module.walk([&](mlir::TFL::WhileOp while_op) {
    auto cond = symbol_table.lookup<mlir::func::FuncOp>(
        while_op->getAttr("cond").cast<mlir::FlatSymbolRefAttr>().getValue());
    AddCallOpInWhileOpRegion(while_op.getCond(), cond);
    while_op->removeAttr("cond");
    auto body = symbol_table.lookup<mlir::func::FuncOp>(
        while_op->getAttr("body").cast<mlir::FlatSymbolRefAttr>().getValue());

        op.getNumOperands(), op.getNumResults()));
  if (op.getCond().front().getNumArguments() !=
      op.getBody().front().getNumArguments())
    return op.emitOpError(llvm::formatv(
        "number of arguments in condition function does not match number of "
        "arguments in body function ({0} != {1})",
        op.getCond().front().getNumArguments(),
        op.getBody().front().getNumArguments()));

        /*parallel_iterations=*/10,
        /*is_stateless=*/false, /*shape_invariant=*/false);
    new_while.getCond().takeBody(while_op.getCond());
    new_while.getBody().takeBody(while_op.getBody());
    ReplaceReturnOp(new_while.getCond(), rewriter);
    ReplaceReturnOp(new_while.getBody(), rewriter);
    rewriter.replaceOp(while_op, new_while.getResults());
    return success();
  }

};

LogicalResult FoldConstantIfOp::matchAndRewrite(
    TF::IfOp op, PatternRewriter& rewriter) const {
  // Extract the constant cond value.
  DenseIntElementsAttr cond_attr;
  if (!matchPattern(op.getCond(), m_Constant(&cond_attr))) return failure();

  // Cond value must be a scalar.
  if (cond_attr.getNumElements() != 1) return failure();

  // Select a branch function.

      return subgraph_index_map_.at(call_op.getCallee().str());
    return std::nullopt;
  };
  auto body_subgraph_index = get_call_index(op.getBody().front());
  auto cond_subgraph_index = get_call_index(op.getCond().front());
  if (!body_subgraph_index || !cond_subgraph_index)
    return op.emitOpError("only single call cond/body while export supported"),
           std::nullopt;
  auto builtin_options =