void getDependentDialects(DialectRegistry &registry) const override {
    RegisterAllTensorFlowDialects(registry);
  }

  void runOnOperation() override;
};

}  // namespace

void LegalizeTFGToTFE::runOnOperation() {
  MLIRContext &context = getContext();
  ModuleOp module = getOperation();

  DenseSet<StringRef> func_symbols;
  for (auto &op : module.getBodyRegion().getOps()) {

                          replication_attr);
  *caller = tf_caller;
  return LogicalResult::success();
}

void EmbeddingSequencingPass::runOnOperation() {
  LOG(INFO) << "EmbeddingSequencingPass::runOnOperation()";
  ModuleOp module = getOperation();

  llvm::SetVector<Operation*> forward_pass_ops;
  llvm::SetVector<Operation*> backward_pass_ops;

  // Constructor used by manually creating the pass.
  explicit PrepareQuantizePass(const quant::QuantizationSpecs& quant_specs)
      : use_quantization_flags_(false), quant_specs_(quant_specs) {}

  void runOnOperation() override;

 private:
  // Set the quantization parameters of the input nodes. These parameters are
  // converted from the user specified input value ranges. The input nodes with

  // processing the function for dynamic range quantization. Therefore, this
  // method preprocess the function to remove all stats ops.
  void removeAllStatsOp(func::FuncOp func);

  void runOnOperation() override;

 private:
  // Keeps track of ops whose inputs cannot be quantized due to not meeting the
  // minimum_elements_for_weights threshold. Prevents emitting duplicate

    populateWithGenerated(pattern_list);
    pattern_list.add<SimplifyBroadcastReshape>(context);
    patterns = std::move(pattern_list);
    return success();
  }

  void runOnOperation() override {
    auto func = getOperation();
    if (failed(applyPatternsAndFoldGreedily(func, patterns)))
      signalPassFailure();
  }

  FrozenRewritePatternSet patterns;
};

}  // namespace

  ReplaceStablehloOpsInMainFunctionWithXlaCallModuleOpsPass(
      const ReplaceStablehloOpsInMainFunctionWithXlaCallModuleOpsPass& other) =
      default;

 private:
  void runOnOperation() override;
};

// Creates a unique stablehlo function name based on op order.
std::string CreateStablehloFunctionName(const int id) {
  return Twine("_stablehlo_main_").concat(std::to_string(id)).str();
}

struct ExecutorIslandCoarseningPass
    : public impl::ExecutorIslandCoarseningPassBase<
          ExecutorIslandCoarseningPass> {
  void runOnOperation() override;
};

void ExecutorIslandCoarseningPass::runOnOperation() {
  // Temporary datastructure to keep operands and results for each island.
  // We define it here to grow and reuse the storage for the duration of the
  // pass.

typedef std::unordered_map<std::string, TF::TPUReplicateMetadataOp> MetadataMap;

struct TPUValidateInputsPass
    : public impl::TPUValidateInputsPassBase<TPUValidateInputsPass> {
  void runOnOperation() override;
};
bool IsTpuRegularOp(Operation* op) {
  static auto* ops = [] {
    llvm::SmallDenseSet<mlir::TypeID, 32>* ops_set =
        new llvm::SmallDenseSet<mlir::TypeID, 32>{

  llvm::StringRef getArgument() const final { return "tfl-pick-subgraphs"; }
  llvm::StringRef getDescription() const final {
    return "Pick the best subgraphs to minimize the overall total costs.";
  }
  void runOnOperation() override;

  std::unordered_map<std::string, std::vector<func::FuncOp>>
  CollectSubgraphFuncs(ModuleOp module);

  void BuildSubgraphs(
      func::FuncOp main_fn,

struct ExtractHeadTailOutsideCompilationPass
    : public impl::ExtractHeadTailOutsideCompilationPassBase<
          ExtractHeadTailOutsideCompilationPass> {
  void runOnOperation() override;
};

void ExtractHeadTailOutsideCompilationPass::runOnOperation() {
  auto& side_effect_analysis = getAnalysis<mlir::TF::SideEffectAnalysis>();
  // Get runtime devices information from the closest parent module.
  auto module = getOperation();