from tensorflow.compiler.mlir.quantization.tensorflow.python import py_function_lib
from tensorflow.compiler.mlir.quantization.tensorflow.python import representative_dataset as rd

# LINT.IfChange(quantize_qat_model)
def quantize_qat_model(
    src_saved_model_path: str,
    dst_saved_model_path: str,
    quantization_options_serialized: bytes,
    *,
    signature_keys: list[str],
    signature_def_map_serialized: dict[str, bytes],

  pybind11_protobuf::ImportNativeProtoCasters();

  m.def(
      // If the function signature changes, likely its corresponding .pyi type
      // hinting should also change.
      // LINT.IfChange
      "quantize_qat_model",
      [](const absl::string_view src_saved_model_path,
         const absl::string_view dst_saved_model_path,
         const QuantizationOptions& quantization_options,

    quant_opts: Quantization options.
    signature_def_map: Signature def key -> SignatureDef mapping.
  """
  logging.info('Running static-range quantization for QAT model.')

  pywrap_quantize_model.quantize_qat_model(
      src_saved_model_path,
      dst_saved_model_path,
      quantization_options_serialized=quant_opts.SerializeToString(),
      signature_keys=list(quant_opts.signature_keys),

inline constexpr absl::string_view kTfQuantPtqDynamicRangeStepName =
    "tf_quant_ptq_dynamic_range";
inline constexpr absl::string_view kTfQuantWeightOnlyStepName =
    "tf_quant_weight_only";

absl::StatusOr<ExportedModel> QuantizeQatModel(
    absl::string_view saved_model_path,
    const std::vector<std::string>& signature_keys,
    const std::unordered_set<std::string>& tags,
    const QuantizationOptions& quantization_options);

      /*unfreeze_constants=*/!quantization_options.freeze_all_variables(),
      function_aliases);
}

}  // namespace

absl::StatusOr<ExportedModel> QuantizeQatModel(
    absl::string_view saved_model_path,
    const std::vector<std::string> &signature_keys,
    const std::unordered_set<std::string> &tags,
    const QuantizationOptions &quantization_options) {