py_function_lib_(ABSL_DIE_IF_NULL(py_function_lib)),  // Crash OK
      src_saved_model_path_(src_saved_model_path),
      function_aliases_(std::move(function_aliases)),
      tags_(std::move(tags)),
      signature_def_map_(std::move(signature_def_map)),
      signature_keys_(std::move(signature_keys)) {}

absl::Status CalibrationComponent::ExportToSavedModel(
    ModuleOp module_op, absl::string_view calibration_data_dir,

  std::string src_saved_model_path_;

  // Function alias mapping for pre-calibrated SavedModel. Used to preserve
  // aliased functions.
  absl::flat_hash_map<FunctionName, FunctionAlias> function_aliases_;

  // Tags to identify the MetaGraphDef to load from a SavedModel.
  const std::unordered_set<std::string> tags_;

  const absl::flat_hash_map<std::string, tensorflow::SignatureDef>
      signature_def_map_;

      return %arg : tensor<1x2xf32>
    }
  )mlir");
  ASSERT_TRUE(module_op);

  absl::flat_hash_map<FunctionName, FunctionAlias> function_aliases;
  UpdateFunctionAliases(function_aliases, *module_op);
  EXPECT_THAT(function_aliases, IsEmpty());
}

TEST_F(UpdateFunctionAliasesTest, AliasUpdatedByMlirFunctionName) {
  // MLIR @main function corresponds to the TF function "main_original".

      CreateMlirContextForQuantization();

  absl::StatusOr<absl::flat_hash_map<FunctionName, FunctionAlias>>
      function_aliases = GetFunctionAliases(saved_model_path, tags);
  if (!function_aliases.ok()) {
    return absl::InternalError(absl::StrCat(
        "Failed to get function alias: ", function_aliases.status().message()));
  }

  absl::StatusOr<mlir::OwningOpRef<mlir::ModuleOp>> module =

      /*asset_file_defs=*/{});
  ASSERT_THAT(exported_model.function_aliases(), SizeIs(2));
  EXPECT_TRUE(exported_model.function_aliases().contains("func1"));
  EXPECT_THAT(exported_model.function_aliases().at("func1"), StrEq("alias1"));
  EXPECT_TRUE(exported_model.function_aliases().contains("func2"));
  EXPECT_THAT(exported_model.function_aliases().at("func2"), StrEq("alias2"));
}

TEST(CreateExportedModelTest, CreateExportedModelWithAddedAssetFileDefs) {

  absl::StatusOr<absl::flat_hash_map<FunctionName, FunctionAlias>>
      function_aliases = GetFunctionAliases(src_saved_model_path, tags);
  if (!function_aliases.ok()) {
    return absl::InternalError(absl::StrCat(
        "Failed to get function alias: ", function_aliases.status().message()));
  }

  TF_ASSIGN_OR_RETURN(
      OwningOpRef<ModuleOp> module,

  absl::flat_hash_map<FunctionName, FunctionAlias> function_aliases(
      meta_graph.meta_info_def().function_aliases().begin(),
      meta_graph.meta_info_def().function_aliases().end());
  return function_aliases;
}

void UpdateFunctionAliases(
    absl::flat_hash_map<FunctionName, FunctionAlias>& function_aliases,
    ModuleOp module_op) {
  absl::flat_hash_set<FunctionName> existing_func_names;

  absl::StatusOr<absl::flat_hash_map<FunctionName, FunctionAlias>>
      function_aliases = GetFunctionAliases(src_saved_model_path, tags);
  if (!function_aliases.ok()) {
    return absl::InternalError(absl::StrCat(
        "Failed to get function alias: ", function_aliases.status().message()));
  }

  TF_ASSIGN_OR_RETURN(
      auto module,

// Updates the function aliases. `module_op` may have different
// function names from the original model, so it re-associates the aliases
// with the new function names. Both the input `function_aliases` and the
// returned value are function name -> alias mappings. `function_aliases` is
// the function alias mapping of the original function. The original function's
// name is retrieved by looking at the "tf._original_func_name" string attribute

    function_aliases: Function name -> function alias mapping.
  """
  loader = saved_model_loader.SavedModelLoader(saved_model_dir)
  meta_graph_def = loader.get_meta_graph_def_from_tags(tags)

  for function_name, function_alias in function_aliases.items():
    meta_graph_def.meta_info_def.function_aliases[function_name] = (
        function_alias
    )