// This is a brief description of the pass.
    return "Apply graph optimizations such as fusing and constant folding to "
           "prepare lifting.";
  }

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

  void runOnOperation() override;

 private:
  Option<OpSet> op_set_{
      *this, "target-opset", llvm::cl::init(OpSet::TF),

struct DecomposeReduceDatasetPass
    : public impl::DecomposeReduceDatasetPassBase<DecomposeReduceDatasetPass> {
  void getDependentDialects(DialectRegistry& registry) const override {
    registry.insert<tf_device::TensorFlowDeviceDialect>();
  }

  void runOnOperation() override;
};

// Create the AnonymousIterator for `reduce_dataset` with `dataset_types` using
// `builder`.

    : public PassWrapper<RaiseToTFOpsPass, OperationPass<func::FuncOp>> {
 public:
  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(RaiseToTFOpsPass)

  void getDependentDialects(DialectRegistry& registry) const override {
    registry.insert<TFRDialect, TF::TensorFlowDialect, scf::SCFDialect,
                    arith::ArithDialect, func::FuncDialect>();
  }

  explicit RaiseToTFOpsPass(std::optional<ModuleOp> tfr_module,

  StringRef getDescription() const final {
    return "Constant folds BroadcastInDimOp before binary elementwise ops";
  }
  void getDependentDialects(::mlir::DialectRegistry &registry) const override {}

  void runOnOperation() override {
    RewritePatternSet patterns(&getContext());
    patterns.add<FoldBroadcastInDimBeforeMulOp>(&getContext());
    patterns.add(ConstantFoldMul);

    const std::vector<std::string>& ordered_output_arrays,
    bool experimental_prune_unreachable_nodes_unconditionally,
    const bool disable_vhlo_to_stablehlo) {
  mlir::DialectRegistry registry;
  registry.insert<mlir::arith::ArithDialect, mlir::func::FuncDialect,
                  mlir::quant::QuantizationDialect,
                  mlir::quantfork::QuantizationForkDialect,

namespace mlir {
namespace TFDevice {
namespace {

struct EmbeddingSequencingPass
    : public ::impl::EmbeddingSequencingPassBase<EmbeddingSequencingPass> {
  void getDependentDialects(mlir::DialectRegistry& registry) const override {
    registry.insert<TF::TensorFlowDialect>();
  }

  void runOnOperation() override;
};

template <typename InputContainer>
std::vector<Type> GetValueTypes(const InputContainer& input) {

  std::string saved_model_mlir_path = tensorflow::GetDataDependencyFilepath(
      "tensorflow/compiler/mlir/tfrt/tests/saved_model/testdata/test.mlir");

  mlir::DialectRegistry registry;
  mlir::RegisterAllTensorFlowDialects(registry);
  mlir::MLIRContext context(registry);
  auto module =
      mlir::parseSourceFile<mlir::ModuleOp>(saved_model_mlir_path, &context);
  ASSERT_TRUE(module);

  std::vector<std::string> inputs;

      "tensorflow/compiler/mlir/tfrt/tests/analysis/testdata/test.mlir");

  mlir::DialectRegistry registry;
  tfrt::RegisterTFRTDialects(registry);
  registry.insert<tfrt::fallback_async::FallbackAsyncDialect>();
  registry.insert<tfrt::fallback_sync::FallbackSyncDialect>();
  mlir::MLIRContext context(registry);
  auto module =
      mlir::parseSourceFile<mlir::ModuleOp>(saved_model_mlir_path, &context);
  ASSERT_TRUE(module);

    const std::vector<std::string>& exported_model_signatures,
    const std::vector<std::string>& tag_names,
    absl::string_view input_arg_shapes_str) {
  mlir::DialectRegistry registry;
  RegisterAllTensorFlowDialects(registry);
  mlir::MLIRContext context(registry);
  context.loadAllAvailableDialects();

  auto module =
      TfToStablehlo(input_path, &context, exported_model_signatures, tag_names,

static void RegisterMlirInputDialects(mlir::DialectRegistry& registry) {
  // TODO(b/259459405): Remove support for stablehlo as an input.
  registry
      .insert<mlir::arith::ArithDialect, mlir::func::FuncDialect,
              mlir::TF::TensorFlowDialect, mlir::stablehlo::StablehloDialect,
              mlir::quant::QuantizationDialect>();
  mlir::func::registerAllExtensions(registry);
}