#include "mlir/Pass/PassManager.h"  // from @llvm-project
#include "tensorflow/compiler/mlir/quantization/stablehlo/cc/component.h"
#include "tensorflow/compiler/mlir/quantization/stablehlo/quantization_config.pb.h"

namespace mlir::quant::stablehlo {

// Performs post-calibration graph transformation as part of post-training
// static-range quantization.
//
// The resulting `ModuleOp` contains quantized StableHLO ops serialized in

#include <optional>

#include "absl/types/span.h"
#include "tensorflow/compiler/mlir/quantization/stablehlo/quantization_config.pb.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/calibrator/calibration_statistics.pb.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/calibrator/calibration_statistics_collector_base.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/quantization_options.pb.h"

namespace tensorflow {

#include "tensorflow/compiler/mlir/quantization/common/attrs_and_constraints.h"  // IWYU pragma: keep - required to use `IsSplatValueEqual`.
#include "tensorflow/compiler/mlir/quantization/tensorflow/passes/passes.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"

namespace mlir::quant {
namespace {

// Applies optimization after quantization.
class OptimizePass

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

// Returns quantization scale specs (fixed output, same scale) for a TF op.

==============================================================================*/
#include "tensorflow/compiler/mlir/quantization/stablehlo/cc/debugger.h"

#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
#include "mlir/Support/LLVM.h"  // from @llvm-project
#include "tensorflow/compiler/mlir/quantization/tensorflow/passes/tf_quant_ops.h"

namespace stablehlo::quantization {

void DisableDebugging(mlir::ModuleOp module_op) {
  module_op.walk(

        {
          "shape": [
            256,
            32,
            32,
            3
          ],
          "quantization": {
          }
        },
        {
          "shape": [
            16,
            3,
            3,
            3
          ],
          "quantization": {
          }
        },
        {
          "shape": [
            0
          ],
        },
        {

#include "tensorflow/compiler/mlir/quantization/common/quantization_lib/quantization_utils.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/quantization_options.pb.h"

namespace mlir::quant::stablehlo {

// Returns StableHLO quantization specs for an op.
std::unique_ptr<OpQuantSpec> GetStableHloOpQuantSpec(Operation* op);

// Returns quantization constraints (ex: fixed output, same scale) given

namespace stablehlo::quantization {

using ::tensorflow::quantization::RepresentativeDatasetFile;

absl::StatusOr<absl::flat_hash_map<std::string, RepresentativeDatasetFile>>
CreateRepresentativeDatasetFileMap(absl::Span<const RepresentativeDatasetConfig>
                                       representative_dataset_configs) {
  absl::flat_hash_map<std::string, RepresentativeDatasetFile>

#include "absl/container/flat_hash_map.h"
#include "absl/status/statusor.h"
#include "absl/types/span.h"
#include "tensorflow/compiler/mlir/quantization/stablehlo/quantization_config.pb.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/quantization_options.pb.h"

namespace stablehlo::quantization {

// Translates a set of `RepresentativeDatsetConfig` to signature key ->
// `RepresentativeDatasetFile` mapping. This is useful when using

#define TENSORFLOW_COMPILER_MLIR_QUANTIZATION_STABLEHLO_QUANTIZE_PASSES_H_

#include "mlir/Pass/PassManager.h"  // from @llvm-project
#include "tensorflow/compiler/mlir/quantization/stablehlo/quantization_options.pb.h"

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