public:
  void runOnOperation() override {
    ModuleOp module = getOperation();
    // Pass an invalid session argument, which is a nullptr.
    if (failed(tf_saved_model::LiftVariables(module, /*session=*/nullptr)))
      signalPassFailure();
  }
};

}  // namespace
}  // namespace tf_saved_model

namespace tf_test {

std::unique_ptr<OperationPass<ModuleOp>> CreateLiftVariablesTestPass() {

// `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 =

std::unique_ptr<OperationPass<ModuleOp>> CreateTransposeCommuteOpsPass();

// Create a pass that legalizes MHLO to TF dialect.
std::unique_ptr<OperationPass<ModuleOp>> CreateLegalizeHloToTfPass();

// Creates a pass which replaces a splat constant tensor with a BroadcastInDim
// op.
std::unique_ptr<OperationPass<ModuleOp>> CreateUnfoldSplatConstantPass();

TEST_F(FindMainFuncOpTest, ReturnsMainFuncOp) {
  constexpr absl::string_view kModuleWithMainFunc = R"mlir(
    module {
      func.func @main() -> () {
        return
      }
    }
  )mlir";

  OwningOpRef<ModuleOp> module_op = ParseModuleOpString(kModuleWithMainFunc);
  EXPECT_THAT(*module_op, NotNull());

  func::FuncOp main_func_op = FindMainFuncOp(*module_op);
  EXPECT_THAT(main_func_op, NotNull());
}

  ];
}

def OptimizeGlobalTensorsPass : Pass<"tf-saved-model-optimize-global-tensors", "ModuleOp"> {
  let summary = "Optimize tf_saved_model.global_tensor's.";

  let constructor = "mlir::tf_saved_model::CreateOptimizeGlobalTensorsPass()";
}

def RemoveVariablesInSessionInitializerPass : Pass<"tf-saved-model-remove-vars-in-session-initializer", "ModuleOp"> {
  let summary = "Remove variables in tf saved model's session initializer.";