StringRef getArgument() const final { return "tfr-decompose"; }

  StringRef getDescription() const final {
    return "Decompose TF ops with the registered composition library.";
  }

  void runOnOperation() override;

 private:
  // Apply canonicalization, mainly constant folding, on the function.
  void ApplyCanonicalization();

  // Rewrite unregistered TF ops to TFR func call ops. Return failure if all the

  StringRef getArgument() const override {
    return "quant-insert-main-function";
  }

  StringRef getDescription() const override {
    return "Inserts the main function to the module.";
  }

  void runOnOperation() override;
};

// Checks if a FuncOp is exported.
bool IsExported(func::FuncOp op) {
  auto exported_names =
      op->getAttrOfType<ArrayAttr>(kTfSavedModelExportedNamesAttr);

    : public impl::LegalizeCompositeToCustomOpPassBase<
          LegalizeCompositeToCustomOpPass> {
  using LegalizeCompositeToCustomOpPassBase::
      LegalizeCompositeToCustomOpPassBase;

  void runOnOperation() override {
    func::FuncOp fn = getOperation();

    int num_layers = 0, current_layer_index = 0;
    // First walk the function to count number of KV Caches.
    fn.walk([&](Operation* op) {

  void getDependentDialects(DialectRegistry& registry) const override {
    registry.insert<mlir::tf_saved_model::TensorFlowSavedModelDialect,
                    ml_program::MLProgramDialect>();
  }
  void runOnOperation() override {
    auto module = getOperation();

    DataFlowSolver solver;
    solver.load<dataflow::DeadCodeAnalysis>();
    solver.load<dataflow::SparseConstantPropagation>();

  void ConvertTFImplementsWithAttributes(func::FuncOp func, FuncAttr attr);
  void ConvertTFAPIImplements(func::FuncOp func, StringAttr attr,
                              ModuleOp module);
  void runOnOperation() override;
};

LogicalResult CheckFusableLayerNormalizedLstmCellSimple(
    func::FuncOp lstm_func) {
  for (int i = 0; i < 5; ++i) {
    auto input = lstm_func.getArgument(i);

    return "quant-prepare-quantize";
  }
  StringRef getDescription() const final {
    // This is a brief description of the pass.
    return "Prepare TF dialect for quantization";
  }

  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

struct ReplicateToIslandPass
    : public impl::ReplicateToIslandPassBase<ReplicateToIslandPass> {
  explicit ReplicateToIslandPass(bool legacy_graph_export) {
    legacy_graph_export_ = legacy_graph_export;
  }
  void runOnOperation() override;
};

// Returns whether op requires `_xla_replica_id` attribute.
bool RequiresReplicaIDAttribute(Operation* op) {
  return llvm::isa<TF::EnqueueTPUEmbeddingSparseTensorBatchOp,

#include "tensorflow/compiler/mlir/tensorflow/transforms/sparsecore/sparsecore_passes.h.inc"

struct EmbeddingProgramKeyPass
    : public impl::EmbeddingProgramKeyPassBase<EmbeddingProgramKeyPass> {
  void runOnOperation() override;
};

// Checks if `op` is nested in `block`.
bool OpInBlock(Operation* op, Block* block) {
  Block* op_block = op->getBlock();
  while (op_block) {
    if (op_block == block) return true;

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

    RewritePatternSet patterns(&getContext());
    patterns

 public:
  StringRef getArgument() const final { return "fuse-mhlo-convolution-pass"; }
  StringRef getDescription() const final {
    return "Fuses MHLO binary element-wise ops and convolution op";
  }

  void runOnOperation() override {
    RewritePatternSet patterns(&getContext());
    patterns.add<FuseMhloMulAndConvolutionPattern>(&getContext());
    if (failed(applyPatternsAndFoldGreedily(getOperation(),