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)) {

             /*use_region_args=*/false,
             /*forward_block_args=*/false);
  CreateCall(if_op, if_op.else_function(),
             /*caller_region=*/if_region.getElseBranch(), if_op.getInput(),
             /*use_region_args=*/false,
             /*forward_block_args=*/false);
  if_op.replaceAllUsesWith(if_region.getResults());
  if_op.erase();
  return success();
}

    (*buffer_to_size)[new_op.getResult(std::get<0>(entry))] = {
        new_op.getResult(std::get<1>(entry)), std::get<2>(entry)};
  }

  new_op.getThenBranch().takeBody(if_op.getThenBranch());
  new_op.getElseBranch().takeBody(if_op.getElseBranch());

  if_op.replaceAllUsesWith(
      new_op.getResults().take_front(if_op.getNumResults()));
  if_op.erase();
  return success();
}

LogicalResult HandleCaseRegionOp(

  // parallelization.
  dataset_if->setAttr("_lower_using_switch_merge", builder.getBoolAttr(true));
  // Empty else branch, if there is no more data, do nothing.
  auto& else_branch = dataset_if.getElseBranch();
  else_branch.push_back(new Block);
  builder.setInsertionPointToEnd(&else_branch.front());
  // Return only the state variables from the body arguments.
  SmallVector<Value, 4> else_returns;

    extern_values.push_back(first);
    return true;
  };

  const TrivialTransformInfo tti(IsSingleCallRegion(if_region.getThenBranch()),
                                 IsSingleCallRegion(if_region.getElseBranch()),
                                 if_arg_matcher);

  std::string then_name, else_name;

  if (tti.can_transform) {
    // We can transform to functional form trivially without outlining.

              /*control_above=*/true, is_map_oc, communication_key_index)))
        return WalkResult::interrupt();
      if (failed(MoveToHostMultiCluster(
              device_cluster, &if_op.getElseBranch().front(),
              {host_if.getElseBranch().front().getTerminator()},
              core_to_compilation_key, core_to_device_ordinal,
              /*control_above=*/true, is_map_oc, communication_key_index)))

    }
  } else if (auto if_op = dyn_cast<TF::IfRegionOp>(user_op)) {
    (*arguments_to_erase)[if_op].push_back(argument_index);
    for (auto callee : {&if_op.getThenBranch(), &if_op.getElseBranch()}) {
      work_list->push_back(std::make_pair(callee, argument_index));
    }
  } else if (auto while_op = dyn_cast<TF::WhileOp>(user_op)) {
    (*arguments_to_erase)[while_op].push_back(argument_index);

LogicalResult IfRegionOp::verifyRegions() {
  IfRegionOp op = *this;
  TypeRange then_types =
      op.getThenBranch().front().getTerminator()->getOperandTypes();
  TypeRange else_types =
      op.getElseBranch().front().getTerminator()->getOperandTypes();

  TypeRangeWithDesc results{op.getResultTypes(), "result"};
  TypeRangeWithDesc then_results{then_types, "then result"};

        /*_else_func_name=*/nullptr);
    new_op.getThenBranch().takeBody(op.getTrueBranch());
    new_op.getElseBranch().takeBody(op.getFalseBranch());
    ReplaceReturnOp(new_op.getThenBranch(), rewriter);
    ReplaceReturnOp(new_op.getElseBranch(), rewriter);
    rewriter.replaceOp(op, new_op.getResults());
    return success();
  }
};

// Converts mhlo.pad to tf.PadV2

  auto opcode_index = GetOpcodeIndex("if", tflite::BuiltinOperator_IF);
  int then_subgraph_index = subgraph_index_map_.at(op.getThenBranch().str());
  int else_subgraph_index = subgraph_index_map_.at(op.getElseBranch().str());
  auto builtin_options = tflite::CreateIfOptions(builder_, then_subgraph_index,
                                                 else_subgraph_index)
                             .Union();