MLIRImportOptions options,
    std::unique_ptr<tensorflow::SavedModelBundle>* saved_model_bundle) {
  // Create local bundle if no one is provided to use.
  std::unique_ptr<tensorflow::SavedModelBundle> bundle;
  if (saved_model_bundle == nullptr) {
    bundle = std::make_unique<tensorflow::SavedModelBundle>();
  } else if (*saved_model_bundle == nullptr) {

                      const RunOptions& run_options, const string& export_dir,
                      const std::unordered_set<string>& tags,
                      SavedModelBundle* const bundle) {
  return LoadSavedModelGeneric<SavedModelBundle>(session_options, run_options,
                                                 export_dir, tags, bundle);
}

Status RestoreSession(const RunOptions& run_options,

    bool include_variables_in_initializers, bool upgrade_legacy,
    bool show_debug_info, TF_Status* status) {
  // Load the saved model into a SavedModelBundle.

  std::unordered_set<string> tag_set =
      absl::StrSplit(tags, ',', absl::SkipEmpty());

  tensorflow::SavedModelBundle bundle;
  auto load_status =
      tensorflow::LoadSavedModel({}, {}, saved_model_path, tag_set, &bundle);
  if (!load_status.ok()) {

namespace odml {

absl::StatusOr<OwningOpRef<mlir::ModuleOp>> ImportSavedModelOrMLIR(
    const std::string& input_path, MLIRContext* context,
    llvm::SourceMgr* source_mgr,
    std::unique_ptr<tensorflow::SavedModelBundle>* saved_model_bundle) {
  if (absl::EndsWith(input_path, ".mlir")) {
    auto file_or_err = llvm::MemoryBuffer::getFileOrSTDIN(input_path.c_str());
    if (std::error_code error = file_or_err.getError()) {

    const absl::string_view saved_model_dir,
    const std::unordered_set<std::string>& saved_model_tags,
    QuantizationConfig* quantization_config,
    const PyFunctionLibrary* quantization_py_function_lib,
    const SavedModelBundle* saved_model_bundle, mlir::PassManager& pass_manager,
    mlir::StatusScopedDiagnosticHandler& status_handler, ModuleOp& module) {
  // TODO: b/194747383 - We need to valid that indeed the "main" func is
  // presented.

    absl::flat_hash_map<std::string, std::string> &function_aliases) {
  // Convert the SavedModelBundle to an MLIR module.
  MLIRImportOptions import_options;
  import_options.upgrade_legacy = true;
  import_options.lift_variables = false;
  import_options.include_variables_in_initializers = true;
  auto bundle = std::make_unique<SavedModelBundle>();

  // TODO: b/213406917 - Add support for the object graph based saved model.

  std::vector<std::string> custom_opdefs(toco_flags.custom_opdefs().begin(),
                                         toco_flags.custom_opdefs().end());
  auto bundle = std::make_unique<tensorflow::SavedModelBundle>();
  TF_ASSIGN_OR_RETURN(
      auto module,
      ImportSavedModel(
          model_flags.saved_model_dir(), model_flags.saved_model_version(),
          tags, absl::MakeSpan(custom_opdefs), exported_names, specs,

    llvm::errs() << "You must specify `emit-select-tf-ops=true` when passing "
                    "`select-user-tf-ops` flag.";
    return kTrFailure;
  }

  std::unique_ptr<tensorflow::SavedModelBundle> bundle;

  // TODO(b/147435528): We need to test the e2e behavior once the graph freezing
  // inside mlir is done.
  if ((import_saved_model_object_graph || import_saved_model_signature_defs) &&

    mlir::OwningOpRef<mlir::ModuleOp> module,
    const mlir::TFL::PassConfig& pass_config,
    const std::unordered_set<std::string>& saved_model_tags,
    std::string* result, SavedModelBundle* saved_model_bundle,
    const PyFunctionLibrary* quantization_py_function_lib) {
  if (toco_flags.has_dump_graphviz_dir()) {
    TF_RETURN_IF_ERROR(DumpOpGraphToFile(
        module.get(),