using ::stablehlo::quantization::QuantizationOptions;

// Returns QuantizationOptions filled with detailed specs when user specifies
// an optional preset method name. The preset methods are defined in
// quantization_options.proto. This function will only be executed if a user
// gives a preset method, not a custom method.
QuantizationOptions FillPresetQuantizationOptions(
    QuantizationOptions quantization_options);

    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 mlir::quant {
namespace {

using QuantizationOptions = tensorflow::quantization::QuantizationOptions;
using QuantizationComponentSpec =
    tensorflow::quantization::QuantizationComponentSpec;

TEST(TfOpQuantSpecTest, WeightComponentSpecExist) {
  QuantizationOptions quant_options;
  QuantizationComponentSpec quant_spec;
  quant_spec.set_quantization_component(

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_ =