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

}

// CPU hardware class which handles CPU capabilities in TFLite.
// This is used by TAC to get op supported/ op cost estimates on CPU.
class CpuHardware : public TargetHardware {
 public:
  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(CpuHardware)

  // String Identifier for CPU hardware.
  static constexpr char kId[] = "CPU";

  double GetHardwareSwitchingCost(const TargetHardware* from,

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

namespace {

// Applies clean-up patterns after quantization.
class PostQuantizePass : public impl::PostQuantizePassBase<PostQuantizePass> {
 public:
  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(PostQuantizePass)

  explicit PostQuantizePass() = default;

 private:
  void runOnOperation() override;
};

// TODO: b/305815328 - Consider preserving leading and trailing QDQs for

  }
  return scheduled_everything;
}

class PartitionedTopologicalSortPass
    : public impl::PartitionedTopologicalSortPassBase<
          PartitionedTopologicalSortPass> {
 public:
  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(PartitionedTopologicalSortPass);

  explicit PartitionedTopologicalSortPass(
      const std::function<bool(Operation *)> &partition = IsFlexDelegate)
      : partition_(partition) {}

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

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

struct SplitMergedOperandsPass
    : public impl::SplitMergedOperandsPassBase<SplitMergedOperandsPass> {
  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(SplitMergedOperandsPass)

  void runOnOperation() override;
};

LogicalResult DuplicateValueIfNeeded(Operation* op,
                                     llvm::DenseSet<Value>* values,