using mlir::Operation;
using mlir::WalkResult;

class VerifyClusteringPass
    : public impl::VerifyClusteringPassBase<VerifyClusteringPass> {
 public:
  void runOnOperation() override;
};

void VerifyClusteringPass::runOnOperation() {
  Operation* func_op = getOperation();

  auto walk_result = func_op->walk([&](Operation* op) {
    if (!tensorflow::tf2xla::internal::IsInBridgeAcceptableDialects(op)) {

  StringRef getDescription() const final {
    return "Replaces operands and results that has quantized type with their "
           "storage types.";
  }
  void runOnOperation() override;
};

void RewriteQuantizedIOPass::runOnOperation() {
  ModuleOp module = getOperation();
  OpBuilder builder(module);
  module.walk([&](func::FuncOp func) {
    Block& block = func.front();

 public:
  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(CompositeLoweringPass);

  void runOnOperation() override;
};

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

void CompositeLoweringPass::runOnOperation() {
  RewritePatternSet patterns(&getContext());

  populateWithGenerated(patterns);

  ConversionTarget target(getContext());

  }
  void getDependentDialects(DialectRegistry& registry) const override {
    registry.insert<TF::TensorFlowDialect>();
  }
  void runOnOperation() override;
};

void DeviceTransformNNAPIPass::runOnOperation() {
  auto func = getOperation();
  auto* ctx = &getContext();
  NNAPIHardware nnapi_hardware;
  RewritePatternSet patterns = nnapi_hardware.GetTransformations(ctx);

constexpr char kDynamicFunctionName[] = "kDynamicFunction";

class InputMetricsLoweringPass
    : public impl::InputLoweringMetricsPassBase<InputMetricsLoweringPass> {
 public:
  void runOnOperation() override;
};

void InputMetricsLoweringPass::runOnOperation() {
  bool has_dynamic_op = false;
  Operation* func_op = getOperation();

  func_op->walk([&](Operation* op) {
    auto abstractOp = op->getRegisteredInfo();

#include "tensorflow/compiler/mlir/tf2xla/transforms/xla_legalize_tf_passes.h.inc"

class InfeedsOpsXlaAdjustLayout
    : public impl::InfeedOpsXlaAdjustLayoutBase<InfeedsOpsXlaAdjustLayout> {
 public:
  void runOnOperation() override;

 private:
  static void runOnInfeedOp(::mlir::mhlo::InfeedOp op) {
    OpBuilder builder(op.getContext());
    SmallVector<Type> result_types(op.getResultTypes().begin(),

struct ConvertLaunchFuncToTFCallPass
    : public impl::ConvertLaunchFuncToTFCallPassBase<
          ConvertLaunchFuncToTFCallPass> {
  void runOnOperation() override;
};

void ConvertLaunchFuncToTFCallPass::runOnOperation() {
  ModuleOp module = getOperation();
  module.walk([&](tf_device::LaunchFuncOp launch) {
    OpBuilder builder(launch);
    auto call_op = builder.create<TF::PartitionedCallOp>(

  FreezeAssetsPass(const FreezeAssetsPass& pass) {}
  explicit FreezeAssetsPass(std::string saved_model_dir) {
    this->saved_model_dir = saved_model_dir;
  }
  void runOnOperation() override;

 private:
  std::string saved_model_dir;
};

void FreezeAssetsPass::runOnOperation() {
  auto module = getOperation();
  if (!tf_saved_model::HasTfSavedModelSemantics(module)) {
    return;
  }
  SymbolTable symbol_table(module);

      const bool enable_per_channel_quantized_weight) {
    enable_per_channel_quantized_weight_ = enable_per_channel_quantized_weight;
  }

 private:
  void runOnOperation() override;
};

void QuantizeCompositeFunctionsPass::runOnOperation() {
  MLIRContext& ctx = getContext();

  PassManager pm(&ctx);
  // Intermediate output from QuantizePass will have quantized ops

#include "tensorflow/compiler/mlir/tensorflow/transforms/tf_passes.h.inc"

struct TPUPartitionedOpConversionPass
    : public impl::TPUPartitionedOpConversionPassBase<
          TPUPartitionedOpConversionPass> {
  void runOnOperation() override;
};

template <typename T>
LogicalResult ReplacePartitionedOp(IntegerAttr num_cores_per_replica, T op) {
  constexpr bool is_input =
      std::is_same_v<std::decay_t<T>, TF::TPUPartitionedInputOp>;