// 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

      TestLiftQuantizableSpotsAsFunctionsWithQuantizationSpecsPass>::
      TestLiftQuantizableSpotsAsFunctionsWithQuantizationSpecsPassBase;

  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(
      TestLiftQuantizableSpotsAsFunctionsWithQuantizationSpecsPass)

 private:
  void runOnOperation() override;
};

// `TestQuantizationSpecs` -> predefined `QuantizationSpecs` textproto.

#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,

// '_from_xla_call_module' trait.
class UnwrapXlaCallModuleOpPass
    : public impl::UnwrapXlaCallModuleOpPassBase<UnwrapXlaCallModuleOpPass> {
 public:
  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(UnwrapXlaCallModuleOpPass)

  explicit UnwrapXlaCallModuleOpPass() = default;

 private:
  void runOnOperation() override;
};

void UnwrapXlaCallModuleOp(TF::XlaCallModuleOp call_op,

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

// The pass to trim functions before we legalize to TFL
// dialect using the specified allowlist.
class TrimFunctionsPass
    : public impl::TrimFunctionsPassBase<TrimFunctionsPass> {
 public:
  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TrimFunctionsPass)

  explicit TrimFunctionsPass() {}
  explicit TrimFunctionsPass(llvm::ArrayRef<std::string> trim_funcs_allowlist) {
    this->trim_funcs_allowlist_ = trim_funcs_allowlist;
  }

          ColocateTPUCopyWithDynamicShapePass> {
  void runOnOperation() override;
};

class Device : public dataflow::AbstractSparseLattice {
 public:
  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(Device)
  using AbstractSparseLattice::AbstractSparseLattice;

  ChangeResult meet(const dataflow::AbstractSparseLattice &other) override {

namespace quant {
namespace {

// Applies all the clean up steps after quantization.
class PostQuantizePass
    : public PassWrapper<PostQuantizePass, OperationPass<func::FuncOp>> {
 public:
  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(PostQuantizePass)

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

  StringRef getArgument() const final {

namespace {
class ConvertXlaCallModuleOpToBfloat16Pass
    : public impl::ConvertXlaCallModuleOpToBfloat16PassBase<
          ConvertXlaCallModuleOpToBfloat16Pass> {
 public:
  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(
      ConvertXlaCallModuleOpToBfloat16Pass)

  explicit ConvertXlaCallModuleOpToBfloat16Pass() = default;

 private:
  void runOnOperation() override;
};