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

void AddQuantizePtqPreCalibrationPasses(
    mlir::OpPassManager &pm, const QuantizationOptions &quantization_options);

void AddQuantizePtqPostCalibrationPasses(
    mlir::OpPassManager &pm, const QuantizationOptions &quantization_options,

namespace py = pybind11;

namespace {

using ::tensorflow::SignatureDef;
using ::tensorflow::quantization::ExportedModel;
using ::tensorflow::quantization::PyFunctionLibrary;
using ::tensorflow::quantization::QuantizationOptions;
using ::tensorflow::quantization::QuantizeDynamicRangePtq;
using ::tensorflow::quantization::QuantizeQatModel;
using ::tensorflow::quantization::QuantizeStaticRangePtq;
using ::tensorflow::quantization::QuantizeWeightOnly;

    const std::unordered_set<std::string>& tags,
    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);

  spec->set_quantization_component(component);
  spec->set_bit_type(bit_type);
  spec->set_bit_width(bit_width);
}

::stablehlo::quantization::QuantizationOptions FillPresetQuantizationOptions(
    ::stablehlo::quantization::QuantizationOptions quantization_options_) {
  CustomQuantizationMethod custom_method =
      quantization_options_.quantization_method().custom_quantization_method();

# quantized results
@test_util.run_all_in_graph_and_eager_modes
class QuantizationOptionsTest(quantize_model_test_base.QuantizedModelTest):
  """Test cases regarding the use of QuantizationOptions proto.

  Run all tests cases in both the graph mode (default in TF1) and the eager mode
  (default in TF2) to ensure support for when TF2 is disabled.
  """

  class SimpleModel(module.Module):

      {asset_file_defs.begin(), asset_file_defs.end()});
}

absl::StatusOr<ExportedModel> ExportCalibrationModel(
    mlir::ModuleOp module_op, mlir::MLIRContext *context,
    const QuantizationOptions &quantization_options,
    const absl::flat_hash_map<std::string, std::string> &function_aliases,
    absl::string_view calibration_data_dir) {
  // Clone ModuleOp and function aliases so changes in this pipeline won't

from tensorflow.python.util import tf_export

# Type aliases for quant_opts_pb2 messages.
_QuantizationOptions = tf_export.tf_export(
    'quantization.experimental.QuantizationOptions'
)(quant_opts_pb2.QuantizationOptions)

_QuantizationMethod = tf_export.tf_export(
    'quantization.experimental.QuantizationMethod'
)(quant_opts_pb2.QuantizationMethod)

_QuantizationComponentSpec = tf_export.tf_export(

namespace stablehlo {
namespace quantization {

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

#include "tensorflow/compiler/mlir/tensorflow/transforms/passes.h"
#include "xla/mlir_hlo/mhlo/transforms/passes.h"

namespace tensorflow {
namespace quantization {
namespace {

using ::tensorflow::quantization::QuantizationOptions;

void AddConvertTpuToCpuModelPasses(mlir::OpPassManager &pm) {
  pm.addPass(mlir::quant::CreateConvertTpuModelToCpuPass());
  pm.addPass(mlir::createInlinerPass());

    const QuantizationComponentSpec expected_bias_component) {
  QuantizationOptions quantization_options;
  quantization_options.mutable_quantization_method()
      ->mutable_preset_quantization_method()
      ->set_preset_method(preset_quantization_options);
  QuantizationOptions filled_quantization_options =
      quant::stablehlo::FillPresetQuantizationOptions(quantization_options);