syntax = "proto3";

package stablehlo.quantization;

option cc_enable_arenas = true;

// Defines arious options to specify and control the behavior of the
// StableHLO quantizer.
// NEXT ID: 2
message QuantizationOptions {
  QuantizationMethod quantization_method = 1;
}

// NEXT ID: 3
message QuantizationMethod {
  // Quantization Method can be either preset or custom.
  oneof quantization_method {

// None, the default configuration is "do not quantize the model".
// 2) A set of supported operations.
// 3) Unit wise quantization precision.
// 4) Target hardware name.
// NEXT ID: 18
message QuantizationOptions {
  // The default quantization configuration for the model. If the below
  // unit-wise configuration does not exist, we use this quantization
  // configuration for the entire model. For each method, default configuration

  Args:
    quantization_options: An instance of QuantizationOptions.
  """
  if quantization_options.op_set == quant_opts_pb2.OpSet.OP_SET_UNSPECIFIED:
    quantization_options.op_set = quant_opts_pb2.OpSet.XLA

  if not quantization_options.tags:
    quantization_options.tags.append(tag_constants.SERVING)

  if not quantization_options.signature_keys:

      ImportAndPreprocessSavedModel(
          saved_model_path,
          {quantization_options.signature_keys().begin(),
           quantization_options.signature_keys().end()},
          {quantization_options.tags().begin(),
           quantization_options.tags().end()},
          context.get(), /*is_inliner_run=*/true,
          /*run_tf_to_stablehlo=*/is_stablehlo,

      quantization_options.quantization_method().preset_method(),
      quantization_options.op_set()));
  pm.addPass(mlir::quant::CreateQuantizeCompositeFunctionsPass(
      quantization_options.quantization_method().preset_method(),
      quantization_options.op_set(),
      quantization_options.enable_per_channel_quantization(),
      quantization_options.min_num_elements_for_weights(),

void AddQuantizePtqDynamicRangePasses(
    mlir::OpPassManager &pm, const QuantizationOptions &quantization_options,
    std::optional<const absl::string_view> mlir_dump_file_prefix =
        std::nullopt);

void AddQuantizeWeightOnlyPasses(
    mlir::OpPassManager &pm, const QuantizationOptions &quantization_options,
    std::optional<const absl::string_view> mlir_dump_file_prefix =
        std::nullopt);

#include "tensorflow/compiler/mlir/quantization/stablehlo/quantization_options.pb.h"
#include "tensorflow/compiler/mlir/quantization/stablehlo/utils/fill_quantization_options.h"

namespace stablehlo {
namespace quantization {

void AddQuantizationPasses(mlir::PassManager& pass_manager,
                           const QuantizationOptions& quantization_options) {
  QuantizationOptions quantization_options_ = quantization_options;

        std::unordered_set<std::string> tags;
        tags.insert(quantization_options.tags().begin(),
                    quantization_options.tags().end());
        const absl::StatusOr<ExportedModel> exported_model = QuantizeQatModel(
            src_saved_model_path, signature_keys, tags, quantization_options);
        if (!exported_model.ok()) return exported_model.status();

    const QuantizationOptions& quantization_options);

// Applies post-training dynamic-range quantization to the model.
absl::StatusOr<ExportedModel> QuantizeDynamicRangePtq(
    absl::string_view saved_model_path,
    const std::vector<std::string>& signature_keys,
    const std::unordered_set<std::string>& tags,
    const QuantizationOptions& quantization_options);

        'Do not specify both the `representative_dataset` argument and the'
        ' `representative_datasets` field in `QuantizationOptions`',
    ):
      quantize_model.quantize(
          self._input_saved_model_path,
          self._output_saved_model_path,
          quantization_options=quantization_options,
          representative_dataset=representative_dataset,
      )

    converted_model = quantize_model.quantize(