#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);

  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(QuantizeWeightsPass)

  explicit QuantizeWeightsPass() : test_mode_(true) { initializeForTest(); }

  explicit QuantizeWeightsPass(
      const tensorflow::quantization::QuantizationOptions& quant_options)
      : test_mode_(false), quant_options_(quant_options) {}

  QuantizeWeightsPass(const QuantizeWeightsPass& other) {
    test_mode_ = other.test_mode_;
    quant_options_ = other.quant_options_;

using ::tensorflow::quantization::OpSet;
using ::tensorflow::quantization::QuantizationComponentSpec;
using ::tensorflow::quantization::QuantizationMethod;
using ::tensorflow::quantization::QuantizationOptions;
using ::tensorflow::quantization::UnitWiseQuantizationSpec;

class LiftQuantizableSpotsAsFunctionsPass
    : public PassWrapper<LiftQuantizableSpotsAsFunctionsPass,
                         OperationPass<ModuleOp>> {

namespace quantization {
// Adds passes for quantization of individual quantizable components.
// (i.e. activation, weight, bias)
void AddQuantizationPasses(mlir::PassManager& pass_manager,
                           const QuantizationOptions& quantization_options);

}  // namespace quantization
}  // namespace stablehlo

    : public internal::SerializedProtobufCaster<
          tensorflow::quantization::ExportedModel> {};

template <>
struct type_caster<tensorflow::quantization::QuantizationOptions>
    : public internal::SerializedProtobufCaster<
          tensorflow::quantization::QuantizationOptions> {};

template <>
struct type_caster<::stablehlo::quantization::CalibrationOptions>
    : public internal::SerializedProtobufCaster<

// bf16 is quantizable.
bool IsValueWithQuantizablePrecision(Value val);

std::optional<tensorflow::quantization::QuantizationComponentSpec>
GetWeightComponentSpec(
    const tensorflow::quantization::QuantizationOptions& quantization_options);

// Returns the spec for the given operation that can be used for both of
// dynamic and static range quantization.
std::unique_ptr<OpQuantSpec> GetTFOpQuantSpec(Operation* op);

          path_map={'serving_default': '/tmp/representative_dataset_path'}
      ).save({'serving_default': representative_dataset})
  )

  # Using in QuantizationOptions.
  quantization_options = tf.quantization.experimental.QuantizationOptions(
      signature_keys=['serving_default'],
      representative_datasets=dataset_file_map,
  )
  tf.quantization.experimental.quantize_saved_model(
      '/tmp/input_model',

    temp_path = self.create_tempdir().full_path
    saved_model_save.save(
        root, temp_path, signatures=root.add.get_concrete_function()
    )

    quantization_options = quant_opts_pb2.QuantizationOptions(
        quantization_method=quant_opts_pb2.QuantizationMethod(
            preset_method=quant_opts_pb2.QuantizationMethod.PresetMethod.METHOD_STATIC_RANGE_INT8
        ),
        tags={tag_constants.SERVING},

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 {