void runOnOperation() override;
};

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

class DropWhileShapeInvariantInDeviceClusterPass
    : public impl::DropWhileShapeInvariantInDeviceClusterPassBase<
          DropWhileShapeInvariantInDeviceClusterPass> {
  void runOnOperation() override;
};

  }

  StringRef getDescription() const final {
    return "generate a temporary file and invoke InitTextFileToImportPass";
  }

 private:
  void runOnOperation() override;
};

void InitTextFileToImportTestPass::runOnOperation() {
  ModuleOp module = getOperation();

  // Create a temporary vocab file.
  int fd;
  SmallString<256> filename;
  std::error_code error_code =

struct IdentifyDilatedConvPass
    : public impl::IdentifyDilatedConvPassBase<IdentifyDilatedConvPass> {
  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(IdentifyDilatedConvPass)
  void runOnOperation() override;
};

void IdentifyDilatedConvPass::runOnOperation() {
  RewritePatternSet patterns(&getContext());
  auto func = getOperation();

  patterns.add<ConvertTFDilatedConvOp<TF::Conv2DOp>,

 public:
  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(UnfuseMhloBatchNormPass)

  explicit UnfuseMhloBatchNormPass() = default;

 private:
  void runOnOperation() override;
};

void UnfuseMhloBatchNormPass::runOnOperation() {
  MLIRContext* ctx = &getContext();
  RewritePatternSet patterns(ctx);
  mhlo::populateUnfuseBatchNormPatterns(ctx, &patterns);

  if (failed(

 public:
  OptimizeIntGraph() = default;
  OptimizeIntGraph(const OptimizeIntGraph &) = default;
  void runOnOperation() override;
};

#include "tensorflow/compiler/mlir/quantization/stablehlo/passes/bridge/optimize.inc"

void OptimizeIntGraph::runOnOperation() {
  RewritePatternSet patterns(&getContext());
  populateWithGenerated(patterns);
  auto func = getOperation();

class InferenceMetricsPass
    : public impl::InferenceMetricsPassBase<InferenceMetricsPass> {
 public:
  void runOnOperation() override;
};

void InferenceMetricsPass::runOnOperation() {
  bool has_tpu_partitioned_call = false;
  ModuleOp module = getOperation();

  for (auto func_op : module.getOps<func::FuncOp>()) {
    func_op->walk(

 public:
  using impl::RemoveShardingCustomCallPassBase<
      RemoveShardingCustomCallPass>::RemoveShardingCustomCallPassBase;

 private:
  void runOnOperation() override;
};

void RemoveShardingCustomCallPass::runOnOperation() {
  func::FuncOp func_op = getOperation();
  MLIRContext& ctx = getContext();

  RewritePatternSet patterns(&ctx);
  populateWithGenerated(patterns);

  PrintPass(const PrintPass& other);
  void runOnOperation() override;

 private:
  llvm::sys::SmartMutex<true> mutex_;
  raw_ostream* os_;
};

PrintPass::PrintPass(raw_ostream* os) {
  if (os) {
    os_ = os;
  } else {
    os_ = &llvm::errs();
  }
}

PrintPass::PrintPass(const PrintPass& other) : PrintPass(other.os_) {}

void PrintPass::runOnOperation() {

struct StripNoinlineAttributePass
    : public impl::StripNoinlineAttributePassBase<StripNoinlineAttributePass> {
 public:
  // void runOnOperation() override;
  void runOnOperation() override {
    // Strip the "tf._noinline" attribute from top-level functions.
    for (auto func_op : getOperation().getOps<func::FuncOp>())
      func_op->removeAttr("tf._noinline");
  }
};

struct GetArithmeticCountPass
    : public impl::GetArithmeticCountPassBase<GetArithmeticCountPass> {
  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(GetArithmeticCountPass)

  void runOnOperation() override;
};

void GetArithmeticCountPass::runOnOperation() {
  auto func = getOperation();
  OpBuilder builder(func);
  func->walk([&](TflArithmeticCountOpInterface arithmetic_count_op) {
    Operation* op = arithmetic_count_op.getOperation();