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

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