public:
  absl::Status CreateMlirModule(std::string module_string = kMlirModuleStr) {
    TF_ASSIGN_OR_RETURN(
        module_, test::GetMlirModuleFromString(module_string, &context_));

    context_.loadAllAvailableDialects();
    return absl::OkStatus();
  }

  absl::StatusOr<FuncOp> GetMain() {
    func::FuncOp main = module_->lookupSymbol<mlir::func::FuncOp>("main");
    if (!main) {

    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,
                    input_arg_shapes_str, /*is_input_mlir_module=*/false);

  if (!module.ok()) {

}

class LoggingHooksTest : public ::testing::Test {
 public:
  LoggingHooksTest() {
    mlir::RegisterCommonToolingDialects(registry_);
    context_.appendDialectRegistry(registry_);
    context_.loadAllAvailableDialects();
    env_ = Env::Default();

    test_group_name_ = "TestGroup";
    test_dir_ = testing::TmpDir();
    setenv("TF_DUMP_GRAPH_PREFIX", test_dir_.c_str(), 1);
  }

                    mlir::TF::TensorFlowDialect, mlir::func::FuncDialect,
                    mlir::shape::ShapeDialect, mlir::TFR::TFRDialect>();
    mlir::MLIRContext ctx(registry);
    ctx.loadAllAvailableDialects();

    llvm::SourceMgr source_mgr = llvm::SourceMgr();
    source_mgr.AddNewSourceBuffer(llvm::MemoryBuffer::getMemBuffer(input),
                                  llvm::SMLoc());

 public:
  SessionClusterTensorflowDialectTest() {
    mlir::RegisterCommonToolingDialects(registry_);
    context_.appendDialectRegistry(registry_);
    context_.loadAllAvailableDialects();
  }

  absl::Status CreateMlirModule(std::string mlir_module_filename) {
    std::string mlir_module_path = TestDataPath() + mlir_module_filename;
    mlir_module_ =

 public:
  TensorflowDialectToExecutorTest() {
    mlir::RegisterCommonToolingDialects(registry_);
    context_.appendDialectRegistry(registry_);
    context_.loadAllAvailableDialects();
  }

  absl::Status CreateMlirModule(std::string mlir_module_filename) {
    std::string mlir_module_path = TestDataPath() + mlir_module_filename;
    mlir_module_ =

 public:
  TensorflowDialectToExecutorTest() {
    mlir::RegisterCommonToolingDialects(registry_);
    context_.appendDialectRegistry(registry_);
    context_.loadAllAvailableDialects();
  }

  absl::Status CreateMlirModule(std::string mlir_module_filename) {
    std::string mlir_module_path = TestDataPath() + mlir_module_filename;
    mlir_module_ =

                                    "TF Saved Model to Stablehlo converter\n");

  mlir::DialectRegistry registry;
  RegisterAllTensorFlowDialects(registry);
  mlir::MLIRContext context(registry);
  context.loadAllAvailableDialects();

  bool is_input_mlir_module = absl::EndsWith(input_path, ".mlir");
  std::vector<std::string> exported_model_signatures_in_vector =
      absl::StrSplit(exported_model_signatures, ',');

 public:
  FunctionClusterTensorflowDialectTest() {
    mlir::RegisterCommonToolingDialects(registry_);
    context_.appendDialectRegistry(registry_);
    context_.loadAllAvailableDialects();
  }

  absl::Status CreateMlirModule(std::string mlir_module_filename) {
    std::string mlir_module_path = TestDataPath() + mlir_module_filename;
    mlir_module_ =

class LowerClusterToRuntimeOpsTest : public ::testing::Test {
 public:
  LowerClusterToRuntimeOpsTest() {
    mlir::RegisterCommonToolingDialects(registry_);
    context_.appendDialectRegistry(registry_);
    context_.loadAllAvailableDialects();

    env_ = Env::Default();
    test_group_name_ = "TestGroup";
    test_dir_ = testing::TmpDir();
    setenv("TF_DUMP_GRAPH_PREFIX", test_dir_.c_str(), /*overwrite=*/1);
  }