ModuleOp module = op->getParentOfType<ModuleOp>();
    SymbolTable symbol_table(module);

    const func::FuncOp ref_func =
        dyn_cast_or_null<func::FuncOp>(symbol_table.lookup(f_attr.getValue()));
    func::FuncOp new_ref_func = dyn_cast<func::FuncOp>(ref_func->clone());
    return symbol_table.insert(new_ref_func);
  }

  LogicalResult match(LiftedOpT op) const override {

  void runOnOperation() override;

 private:
  LogicalResult ConvertIfOp(SymbolTableCollection& symbol_table,
                            IfRegionOp if_region);
  LogicalResult ConvertCaseOp(SymbolTableCollection& symbol_table,
                              CaseRegionOp case_region);
  LogicalResult ConvertWhileOp(SymbolTableCollection& symbol_table,
                               WhileRegionOp while_region);

  while (!func_ops.empty()) {
    auto main_func = func_ops.back();
    func_ops.pop_back();
    if (!main_func) continue;

    SymbolTable symbol_table(module_op);
    for (auto call_op : main_func.getOps<TF::PartitionedCallOp>()) {
      func_ops.push_back(dyn_cast_or_null<func::FuncOp>(symbol_table.lookup(
          mlir::cast<FlatSymbolRefAttr>(call_op.getFAttr()).getValue())));
    }

  module_for_func.get().getOperation()->setAttr(kVersionsAttr, versions_attr);
  SymbolTable symbol_table(module_for_func.get());

  while (!referenced.empty()) {
    auto func = referenced.pop_back_val();

    // Skip functions that have already been cloned into new module.
    if (symbol_table.lookup<func::FuncOp>(func.getName())) continue;

                               const StringRef func_name) {
  static tensorflow::mutex* mtx = new tensorflow::mutex();
  tensorflow::mutex_lock lock(*mtx);

  SymbolTable symbol_table(&module);
  std::string unique_name = func_name.str();
  int32_t uniquing_counter = 0;
  while (symbol_table.lookup(unique_name) != nullptr) {
    ++uniquing_counter;
    unique_name = absl::StrCat(func_name.str(), "_", uniquing_counter);
  }

                              StringRef error_message) const {
    ModuleOp module = call_op->getParentOfType<ModuleOp>();
    SymbolTable symbol_table(module);
    mlir::func::FuncOp composite_func =
        dyn_cast<func::FuncOp>(symbol_table.lookup(function_name));
    if (!composite_func) return;

    composite_func.walk([&](Operation* op) {
      if (op->hasAttr(kAttrMapAttribute)) {

      builder.getStringAttr(kTfSavedModelInitializerRestoreType));

  builder.setInsertionPointToStart(&init_func.front());
  builder.create<func::ReturnOp>(loc, /*operands=*/ValueRange{});

  SymbolTable symbol_table(module_op);
  symbol_table.insert(init_func);

  AddSymbolToInitializersAttr(
      session_init_op, FlatSymbolRefAttr::get(init_func.getSymNameAttr()));

  return init_func;
}

  if (!main_graph_op) return;

  SessionInitializerOp session_init_op = GetSessionInitializerOp(module_op);
  if (!session_init_op) return;

  // initializer_type -> init_func_op mapping.
  SymbolTable symbol_table{module_op};
  FailureOr<absl::flat_hash_map<std::string, func::FuncOp>> init_func_ops =
      GetInitFuncOps(module_op);
  if (failed(init_func_ops)) {

void CreateFunctions(ModuleOp module_op,
                     llvm::StringMap<FunctionMetadata> &metadatas) {
  MLIRContext *context = module_op.getContext();
  SymbolTable symbol_table(module_op);
  for (auto &iter : metadatas) {
    llvm::StringRef host = iter.first();
    FunctionMetadata &metadata = iter.second;

    // Do not create any new function for the operations on the localhost.

          .getDefiningOp();
  const auto entry_function_symbol_ref =
      lifted_op->getAttrOfType<FlatSymbolRefAttr>("_entry_function");
  SymbolTable symbol_table(*module_op);
  auto entry_func = dyn_cast_or_null<func::FuncOp>(
      symbol_table.lookup(entry_function_symbol_ref.getValue()));
  auto lifted_dot_general_op =
      FindOperationOfType<mlir::stablehlo::DotGeneralOp>(entry_func);