// Convert a TPU model to be compatible on CPU by rewriting/removing TPU ops.
class ConvertTpuModelToCpuPass
    : public PassWrapper<ConvertTpuModelToCpuPass, OperationPass<ModuleOp>> {
 public:
  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(ConvertTpuModelToCpuPass)
  explicit ConvertTpuModelToCpuPass() = default;

  StringRef getArgument() const final {
    // This is the argument used to refer to the pass in

#include "tensorflow/compiler/mlir/lite/stablehlo/transforms/passes.h.inc"

class LegalizeHloToTfLitePass
    : public impl::LegalizeHloToTfLitePassBase<LegalizeHloToTfLitePass> {
 public:
  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(LegalizeHloToTfLitePass);

  void runOnOperation() override;
};

class ConvertReduceOpToTFLiteArgmax
    : public ConvertReduceOpToArgMinMax<TFL::ReduceMaxOp, TFL::ArgMaxOp,

namespace {

class ConvertCustomAggregationOpToQuantStatsPass
    : public PassWrapper<ConvertCustomAggregationOpToQuantStatsPass,
                         OperationPass<func::FuncOp>> {
 public:
  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(
      ConvertCustomAggregationOpToQuantStatsPass)

  StringRef getArgument() const final {
    // This is the argument used to refer to the pass in the textual format (on

namespace {

class MergeFusionWithDequantizePass
    : public impl::MergeFusionWithDequantizePassBase<
          MergeFusionWithDequantizePass> {
 public:
  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(MergeFusionWithDequantizePass)

  explicit MergeFusionWithDequantizePass() = default;

 private:
  void runOnOperation() override;
};

class MergeFusionWithUniformDequantizePattern

#include "tensorflow/compiler/mlir/lite/transforms/passes.h.inc"

struct LegalizeJaxRandomPass
    : public impl::LegalizeJaxRandomPassBase<LegalizeJaxRandomPass> {
 public:
  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(LegalizeJaxRandomPass)

  void runOnOperation() override;
};

inline ConstBytesAttr CustomOption(ImplicitLocOpBuilder *builder,
                                   const std::string &content) {

namespace {

constexpr StringRef kDequantizeFunctionName = "composite_dequantize";

class PropagateQuantizeType
    : public PassWrapper<PropagateQuantizeType, OperationPass<ModuleOp>> {
 public:
  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(PropagateQuantizeType)

  // Constructor used by the PassRegistration. This will remove the adaptor ops.
  explicit PropagateQuantizeType() = default;

  StringRef getArgument() const final {

      total_element_count += input_type.getNumElements();
    }
  }
  return total_element_count;
}

struct GetOpCostPass
    : mlir::PassWrapper<GetOpCostPass, OperationPass<func::FuncOp>> {
  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(GetOpCostPass)

  llvm::StringRef getArgument() const final { return "tfl-get-op-cost"; }
  llvm::StringRef getDescription() const final {
    return "Get cost for every op";
  }

// Adds QuantizationUnitLoc to quantizable layers.
class AddQuantizationUnitLocPass
    : public PassWrapper<AddQuantizationUnitLocPass,
                         OperationPass<func::FuncOp>> {
 public:
  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(AddQuantizationUnitLocPass)
  explicit AddQuantizationUnitLocPass() = default;

  StringRef getArgument() const final {
    // This is the argument used to refer to the pass in

// TensorFlow Lite hash table kernels.
class LegalizeHashTablesPass
    : public impl::LegalizeHashTablesPassBase<LegalizeHashTablesPass> {
 public:
  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(LegalizeHashTablesPass)

  void runOnOperation() override {
    auto module = getOperation();

    if (!checkWhetherGraphHasValidStaticLookupTables(module)) {
      return;
    }

#include "tensorflow/compiler/mlir/lite/transforms/passes.h.inc"

// Applies all the clean up steps after quantization.
class PostQuantizePass : public impl::PostQuantizePassBase<PostQuantizePass> {
 public:
  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(PostQuantizePass)

  // Constructor used by the PassRegistration. This will remove the adaptor ops.
  explicit PostQuantizePass() { this->emit_quant_adaptor_ops_ = false; }