namespace {

using QuantMethod =
    ::tensorflow::quantization::QuantizationMethod::PresetMethod;
using QuantizationUnit = std::pair<Operation*, int>;
using QuantizationUnits = llvm::SetVector<QuantizationUnit>;
using ::tensorflow::quantization::OpSet;

// Preprocesses ops to allow multi-axis quantization, prior to quantization
// passes. Currently, per-channel quantization only supports 1D results.
class PreprocessOpPass

      llvm::cl::values(
          clEnumValN(tensorflow::quantization::QuantizationMethod::
                         METHOD_STATIC_RANGE_INT8,
                     "ptq", "Post-training static-range quantization"),
          clEnumValN(tensorflow::quantization::QuantizationMethod::
                         METHOD_DYNAMIC_RANGE_INT8,
                     "drq", "Post-training dynamic-range quantizaiton"),

    CustomQuantizationMethod custom_quantization_method = 2;
  }
}

// Preset model quantization method for optimization.
//
// Common quantization methods are defined as preset methods in this message.
// Note that the quantization specs (ex: bit width) for preset quantization
// methods are fixed. To use a different quantization spec for a given method,
// use CustomQuantizationMethod.
// NEXT ID: 2

        ":types",
        "//tensorflow/compiler/mlir/quantization/stablehlo:quantization_config_proto_cc",
        "//tensorflow/compiler/mlir/quantization/tensorflow:exported_model_proto_cc",
        "//tensorflow/compiler/mlir/quantization/tensorflow:passes",
        "//tensorflow/compiler/mlir/quantization/tensorflow/cc:convert_asset_args",
        "//tensorflow/compiler/mlir/quantization/tensorflow/cc:run_passes",

        "//tensorflow/compiler/mlir/quantization/stablehlo/cc:config",
        "//tensorflow/compiler/mlir/quantization/stablehlo/cc:static_range_ptq",
        "//tensorflow/compiler/mlir/quantization/stablehlo/cc:weight_only_ptq",
        "//tensorflow/compiler/mlir/quantization/tensorflow/python:py_function_lib",
        "//tensorflow/compiler/mlir/tensorflow/transforms:tensorflow_passes",

  // Defines calibration options for quantization. This option is only used for
  // activation of static range quantization (SRQ). Quantization calibration
  // method is set to MIN_MAX by default.
  stablehlo.quantization.CalibrationOptions calibration_options = 15;

  // Configuration related to quantization debugger.
  stablehlo.quantization.DebuggerConfig debugger_config = 16;

  reserved 3;

)

pytype_strict_library(
    name = "quantization",
    srcs = ["quantization.py"],
    visibility = ["//visibility:public"],
    deps = [
        ":pywrap_quantization",
        "//tensorflow/compiler/mlir/quantization/stablehlo:quantization_config_proto_py",
        "//tensorflow/compiler/mlir/quantization/tensorflow/python:py_function_lib_py",
        "//tensorflow/compiler/mlir/quantization/tensorflow/python:save_model",

    const ::stablehlo::quantization::QuantizationSpecs& specs,
    const ::stablehlo::quantization::DebuggerConfig& debugger_config);

// Adds passes for static-range quantization post-calibration. Utilizes tensor
// statistics collected from the calibration step and performs quantization.
void AddPostCalibrationPasses(
    OpPassManager& pm,

    )

  # In addition to preset quantization methods, fine-grained control of
  # quantization for each component is also supported.
  _QuantizationComponentSpec = (
      tf.quantization.experimental.QuantizationComponentSpec
  )
  quantization_options = tf.quantization.experimental.QuantizationOptions(
      signature_keys=['your_signature_key'],
      quantization_method=tf.quantization.experimental.QuantizationMethod(

#include "tensorflow/compiler/mlir/quantization/stablehlo/instrumentations/save_report.h"
#include "tensorflow/compiler/mlir/quantization/stablehlo/quantization_config.pb.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/cc/run_passes.h"
#include "xla/mlir_hlo/mhlo/transforms/passes.h"
#include "tsl/platform/errors.h"

namespace mlir::quant::stablehlo {

using ::stablehlo::quantization::GetReportFilePath;