PropagatePotentiallyWrittenUpFromCallee(if_op.getThenBranch(),
                                              if_op.getODSOperands(1));
      PropagatePotentiallyWrittenUpFromCallee(if_op.getElseBranch(),
                                              if_op.getODSOperands(1));
      return;
    }
    if (auto while_op = dyn_cast<TF::WhileOp>(op)) {

        rewriter.create<scf::IfOp>(loc, tf_if_region_return_type,
                                   scf_if_condition, /*withElseRegion=*/true);

    Region& then_region = op.getThenBranch();
    Region& else_region = op.getElseBranch();

    // Create the `then` and `else` regions of the `scf.if` op.
    createScfThenOrElse(then_region, scf_if_op.getThenRegion(),
                        tf_if_region_return_type, rewriter);