auto new_op = CreateTfOp<CaseOrIfRegionOp>(rewriter, op, new_result_types,
                                               op.getOperand(), op->getAttrs(),
                                               op.getNumRegions());

    int next_index = 0;
    for (auto result : op.getResults()) {
      if (!result_to_extern_value.count(result)) {
        result.replaceAllUsesWith(new_op.getResult(next_index++));
        continue;

  LogicalResult matchAndRewrite(
      mhlo::WhileOp while_op, OpAdaptor adaptor,
      ConversionPatternRewriter& rewriter) const final {
    // HLO WhileOp should have two regions: cond and body.
    if (while_op->getNumRegions() != 2) return failure();

    // This rule doesn't support mhlo::WhileOp with tuple inputs.
    for (auto type : while_op->getOperandTypes()) {
      if (mlir::isa<TupleType>(type)) return failure();
    }