StringPiece func_name);

  // Runs the decompose passes on the user_module.
  Status DecomposeGraph(mlir::ModuleOp user_module);

  // Erases the tfr_module created.
  void Destroy();

 private:
  mlir::ModuleOp tfr_module_;
  mlir::PassManager pm_;

  GraphExportConfig export_confs_;
};

// Decomposes the NodeDef to a set of primitive ops according to the decompose

    "mlir::TFL::TFLDialect",
  ];
}

def LegalizeHloToTfPass : Pass<"tf-legalize-hlo", "ModuleOp"> {
  let summary = "Legalize from MHLO to the TF dialect";
  let dependentDialects = ["TF::TensorFlowDialect"];
  let constructor = "mlir::odml::CreateLegalizeHloToTfPass()";
}

def LegalizeHloToTfLitePass : Pass<"tfl-legalize-hlo", "mlir::ModuleOp"> {
  let summary = "Legalize from MHLO to the TFLite dialect";

std::unique_ptr<OperationPass<ModuleOp>> CreateQuantizeWeightsPass(
    const tensorflow::quantization::QuantizationOptions& quant_options);

// Propagate quantized type through allowed ops.
std::unique_ptr<OperationPass<ModuleOp>> CreatePropagateQuantizeTypePass();

// Create a pass that inserts dump tensor to quantizable layer's output.
std::unique_ptr<OperationPass<ModuleOp>> CreateAddDumpTensorOpPass(

SessionInitializerOp GetSessionInitializerOp(ModuleOp module_op);

// Returns the exported name for the session initializer function.
SmallVector<StringRef, 2> GetSessionInitializerExportedName(ModuleOp module_op);

// Returns initializer function ops. These functions' symbols are in the
// "initializers" attribute of the session initializer op.
SmallVector<func::FuncOp, 2> GetInitializerFunctions(ModuleOp module_op);

// `module_op` is the input graph to be quantized and it should contain
// StableHLO ops.
//
// Returns a quantized `ModuleOp` in StableHLO, potentially wrapped inside a
// XlaCallModuleOp. Returns a non-OK status if quantization fails, or any of
// `saved_model_bundle` or `quantization_py_function_lib` is a nullptr.
absl::StatusOr<mlir::ModuleOp> RunQuantization(
    const SavedModelBundle* saved_model_bundle,
    absl::string_view saved_model_dir,

}

// Creates a pruned module containing the XlaCallModule's entry function and
// other functions transitively called by the entry function.
FailureOr<OwningOpRef<ModuleOp>> PruneStablehloModule(
    SymbolTableCollection& symbol_table, ModuleOp module, XlaCallModuleOp op) {
  auto entry_func_symbol =
      op->getAttrOfType<FlatSymbolRefAttr>(kStablehloEntryFunctionAttrName);
  if (!entry_func_symbol) {

std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>>
CreateExtractOutsideCompilationPass();

// Create a pass that encapsulates StatefulPartitionedCallOp within a cluster.
std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>>
CreateXlaClusterFormationPass();

// Creates a pass that marks unsupported ops in device cluster for outside
// compilation.
std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>>
CreateMarkOpsForOutsideCompilationPass();

using UpdateFunctionAliasesTest = ::mlir::quant::QuantizationTestBase;

TEST_F(UpdateFunctionAliasesTest, NoAliasesReturnsEmptyMap) {
  // MLIR @main function corresponds to the TF function "main_original".
  OwningOpRef<ModuleOp> module_op = ParseModuleOpString(R"mlir(
    func.func private @main(%arg: tensor<1x2xf32>) -> (tensor<1x2xf32>) attributes {tf._original_func_name = "main_original"} {
      return %arg : tensor<1x2xf32>
    }
  )mlir");

  }

  bool ModulesEqual(const ModuleOp& module_before,
                    const ModuleOp& module_after) {
    return mlir::OperationEquivalence::isEquivalentTo(
        module_before, module_after, mlir::OperationEquivalence::None);
  }

  mlir::LogicalResult Run() {
    mlir::OwningOpRef<mlir::ModuleOp> module_before = module_->clone();
    LogicalResult run_result = pm_->run(module_.get());

  mlir::OwningOpRef<mlir::ModuleOp> module_ref =
      mlir::ModuleOp::create(mlir::UnknownLoc::get(&context));
  unsetenv("TF_DUMP_GRAPH_PREFIX");

  std::string filepath = DumpMlirOpToFile("module", module_ref.get());
  EXPECT_EQ(filepath, "(TF_DUMP_GRAPH_PREFIX not specified)");
}

TEST(DumpMlirModuleTest, LogInfo) {
  mlir::MLIRContext context;
  mlir::OwningOpRef<mlir::ModuleOp> module_ref =