init_func_op.emitError("expects exactly one block in the MLIR function");
      return signalPassFailure();
    }

    auto var_handle_ops =
        init_func_op.getBlocks().front().getOps<VarHandleOp>();
    llvm::SmallVector<VarHandleOp, 4> init_vars(var_handle_ops.begin(),
                                                var_handle_ops.end());
    RemoveVariables(init_vars);
  }
}

}  // namespace

      }
    }
  }
}

void AddFunctionsForExportedNamesPass::runOnOperation() {
  Block& module_body = getOperation().getRegion().front();
  for (auto f :
       llvm::make_early_inc_range(getOperation().getOps<func::FuncOp>())) {
    auto exported_names = GetExportedNames(f);
    if (exported_names.empty() || !f.isPublic()) continue;

    if (exported_names.size() == 1 && exported_names[0] == f.getName()) {

  if (!bundle_writer.status().ok()) {
    return bundle_writer.status();
  }

  std::vector<std::string> saved_variable_shared_names;
  for (auto assign_variable_op :
       session_init_func_type_restore_op.getOps<mlir::TF::AssignVariableOp>()) {
    if (const absl::StatusOr<std::string> variable_shared_name =
            AddTensorToBundleWriter(assign_variable_op, bundle_writer);
        !variable_shared_name.ok()) {

      return failure();

    do {
      SymbolUserMap symbol_users(symbol_table_collection, module);

      made_changes = false;
      for (auto func :
           llvm::make_early_inc_range(module.getOps<func::FuncOp>())) {
        ArrayRef<Operation *> users = symbol_users.getUsers(func);
        if (users.size() <= 1) {
          continue;
        }

  auto func_ops = module_op.getOps<func::FuncOp>();
  auto main_func_itr = absl::c_find_if(func_ops, [&main_func_id](auto func_op) {
    return func_op.getName() == main_func_id;
  });

  if (main_func_itr == func_ops.end()) return {};
  return *main_func_itr;
}

func::FuncOp GetSaveFuncOp(ModuleOp module_op) {
  for (auto func_op : module_op.getOps<func::FuncOp>()) {

  // functions and avoid converting ConstOps inside initializer functions.
  for (auto func_op : module_op.getOps<func::FuncOp>()) {
    // Do not unfreeze constants under these functions.
    if (func_op.getSymName().contains("while_body")) continue;
    if (func_op.getSymName().contains("while_cond")) continue;
    absl::c_copy_if(func_op.getOps<TF::ConstOp>(),
                    std::back_inserter(target_const_ops),

    // input token in the parent op or a new token for the entry computation.
    auto token = StringAttr::get(ctx, tensorflow::kXlaTokenArgNodeName);
    for (Operation& op : region->getOps()) {
      // Only communication related ops that needs to have token should have the
      // extra attribute.
      if (!ops_with_tokens.contains(&op)) continue;

  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(LegalizeWhilePass)
  void RunOnFunction(func::FuncOp func);

  void runOnOperation() override {
    for (auto op : getOperation().getOps<func::FuncOp>()) RunOnFunction(op);
  }
};

}  // namespace

// Inserts call to the given function into the 'region'.
void CreateRegionWithCall(func::FuncOp func, Region& region, Location loc) {

        device_mgr_.get(), Env::Default(), /*config=*/nullptr,
        TF_GRAPH_DEF_VERSION, lib_def_.get(), opts,
        /*default_thread_pool=*/nullptr, /*cluster_flr=*/nullptr);
    flr_ = pflr_->GetFLR("/job:localhost/replica:0/task:0/cpu:0");
  }

  FunctionLibraryRuntime* flr_;
  std::unique_ptr<DeviceMgr> device_mgr_;
  std::unique_ptr<FunctionLibraryDefinition> lib_def_;

            .append("duplicate exported name '", exported_name, "'")
            .attachNote(p.first->getSecond()->getLoc())
            .append("previously seen here");
      }
    }
  }
  for (auto func : module.getOps<func::FuncOp>()) {
    const bool is_exported = IsExported(func);

    if (is_exported && !func.isPublic()) {
      return func.emitError()
             << "exported function @" << func.getName() << " should be public";