absl::StatusOr<OwningOpRef<ModuleOp>> GetMlirModuleFromString(
    absl::string_view module_string, MLIRContext* context) {
  DialectRegistry mlir_registry;
  RegisterCommonToolingDialects(mlir_registry);
  context->appendDialectRegistry(mlir_registry);

  OwningOpRef<ModuleOp> mlir_module;
  auto status =
      tensorflow::DeserializeMlirModule(module_string, context, &mlir_module);
  if (!status.ok()) {
    return status;
  }

PYBIND11_MODULE(mlir_wrapper, m) {
  m.def("preloadTensorFlowDialects", [](mlir::MLIRContext &context) {
    mlir::DialectRegistry registry;
    mlir::RegisterAllTensorFlowDialects(registry);
    context.appendDialectRegistry(registry);
    context.loadAllAvailableDialects();
  });

  m.def("verify", [](std::string input) {
    llvm::SourceMgr SM = llvm::SourceMgr();
    SM.AddNewSourceBuffer(llvm::MemoryBuffer::getMemBuffer(input),

      "tensorflow/compiler/mlir/tf2xla/internal/testdata/");
}

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

}

class SessionClusterTensorflowDialectTest : public ::testing::Test {
 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;

}

class TensorflowDialectToExecutorTest : public ::testing::Test {
 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;

absl::StatusOr<OwningOpRef<ModuleOp>> GetMlirModuleFromString(
    absl::string_view module_string, MLIRContext* context) {
  mlir::DialectRegistry mlir_registry;
  RegisterAllTensorFlowDialects(mlir_registry);
  context->appendDialectRegistry(mlir_registry);

  OwningOpRef<ModuleOp> mlir_module;
  auto status =
      tensorflow::DeserializeMlirModule(module_string, context, &mlir_module);
  if (!status.ok()) {
    return status;
  }

  }
  return count;
}

class TensorflowDialectToExecutorTest : public ::testing::Test {
 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;

}

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

  absl::StatusOr<std::vector<mlir::tf_device::ClusterFuncOp>> GetClusterFuncOps(
      absl::string_view mlir_module_filename) {

  auto module_status = importer_->Import();
  if (!module_status.ok()) {
    return module_status.status();
  }
  auto module = module_status->get();
  auto* context = module->getContext();
  context->appendDialectRegistry(registry_);
  context->loadAllAvailableDialects();

  // Run TAC passes.
  auto status = RunTacPasses(&module, options_.debug_mode);

  if (!status.ok()) {
    return status;
  }

namespace {

std::unique_ptr<MLIRContext> CreateContext() {
  auto context = std::make_unique<MLIRContext>();
  DialectRegistry mlir_registry;
  RegisterCommonToolingDialects(mlir_registry);
  context->appendDialectRegistry(mlir_registry);
  return context;
}

TEST(IsLargeFloatTypeTest, scalars) {
  auto context = CreateContext();

  EXPECT_FALSE(IsLargeFloatType(Float8E4M3FNType::get(context.get())));