}

TEST_F(LoaderTest, RestoreSession) {
  SavedModelBundle bundle;
  SessionOptions session_options;
  RunOptions run_options;

  const string export_dir =
      io::JoinPath(testing::TensorFlowSrcRoot(), kTestDataSharded);
  TF_ASSERT_OK(LoadSavedModel(session_options, run_options, export_dir,
                              {kSavedModelTagServe}, &bundle));

  SavedModelBundle actual_bundle;

  void AddSignatureDefToSavedModelBundle(const SignatureDef& signature_def,
                                         const string& key,
                                         SavedModelBundle* saved_model_bundle) {
    MetaGraphDef* meta_graph_def = &saved_model_bundle->meta_graph_def;
    (*meta_graph_def->mutable_signature_def())[key] = signature_def;
  }

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

    }
  } else {
    tensor_names->insert(tensor_info.name());
  }
}

// Gets the union of all inputs and outputs of all SignatureDefs in the bundle
void GetSignatureDefsInputsAndOutputs(
    const SavedModelBundle& saved_model_bundle,
    std::unordered_set<string>* inputs, std::unordered_set<string>* outputs) {
  for (auto& sigdef_elem : saved_model_bundle.meta_graph_def.signature_def()) {

                      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,