# limitations under the License.
# ==============================================================================
from typing import Any, Optional

class PyFunctionLibrary:

  # LINT.IfChange(save_exported_model)
  def save_exported_model(
      self,
      dst_saved_model_path: str,
      exported_model_serialized: bytes,
      src_saved_model_path: str,
      tags: set[str],
      serialized_signature_def_map: dict[str, bytes],

}  // namespace

PYBIND11_MODULE(pywrap_function_lib, m) {
  py::class_<PyFunctionLibrary, PyFunctionLibraryTrampoline>(
      m, "PyFunctionLibrary")
      .def(py::init<>())
      .def("save_exported_model", &PyFunctionLibrary::SaveExportedModel,
           py::arg("dst_saved_model_path"),
           py::arg("exported_model_serialized"),
           py::arg("src_saved_model_path"), py::arg("tags"),

  //
  // If the function signature changes, likely its corresponding .pyi type
  // hinting and definition should also change.
  // LINT.IfChange(save_exported_model)
  virtual std::optional<bool> SaveExportedModel(
      absl::string_view dst_saved_model_path,
      const ExportedModel& exported_model,
      absl::string_view src_saved_model_path,
      const std::unordered_set<std::string>& tags,

  """

  # LINT.IfChange(save_exported_model)
  def save_exported_model(
      self,
      dst_saved_model_path: str,
      exported_model_serialized: bytes,
      src_saved_model_path: str,
      tags: set[str],
      serialized_signature_def_map: dict[str, bytes],
  ) -> Optional[bool]:
    # LINT.ThenChange(py_function_lib.h:save_exported_model)

        // tensorflow/python/saved_model/tag_constants.py.
        if (quantization_options.has_debugger_config()) {
          tags.erase("tpu");
        }
        py_function_library.SaveExportedModel(
            dst_saved_model_path, *exported_model, src_saved_model_path, tags,
            signature_def_map);

        return absl::OkStatus();
      },
      R"pbdoc(

                          *ctx, *module));

  // Remove the `tpu` tag for exporting because the output quantized model is
  // essentially a CPU model.
  tags.erase("tpu");

  py_function_library.SaveExportedModel(
      dst_saved_model_path, post_calibrated_exported_model,
      src_saved_model_path, tags, signature_def_map);

  return absl::OkStatus();
}

                          *ctx, *module));

  // Remove the `tpu` tag for exporting because the output quantized model is
  // essentially a CPU model.
  tags.erase("tpu");

  py_function_library.SaveExportedModel(
      dst_saved_model_path, post_calibrated_exported_model,
      src_saved_model_path, tags, signature_def_map);

  return absl::OkStatus();
}

                          *cloned_module_ref, checkpoint_dir, function_aliases_,
                          {asset_file_defs.begin(), asset_file_defs.end()}));

  py_function_lib_->SaveExportedModel(dst_saved_model_path, exported_model,
                                      src_saved_model_path_, tags_,
                                      signature_def_map_);

  return absl::OkStatus();
}

  if (calibration_exported_model.has_graph_def()) {
    TF_ASSIGN_OR_RETURN(std::string calibration_saved_model_dir,
                        CreateTmpDir());
    py_function_library.SaveExportedModel(
        calibration_saved_model_dir, calibration_exported_model,
        saved_model_path, tags, signature_def_map);

    py_function_library.RunCalibration(