mlir::MLIRContext mlir_context(mlir_registry);
  mlir::OwningOpRef<mlir::ModuleOp> mlir_module_with_add, mlir_module_with_sub;
  TF_ASSERT_OK(DeserializeMlirModule(module_with_add, &mlir_context,
                                     &mlir_module_with_add));
  TF_ASSERT_OK(DeserializeMlirModule(module_with_sub, &mlir_context,
                                     &mlir_module_with_sub));

std::string SerializeMlirModule(mlir::ModuleOp module_op);

// Parses a MLIR module from `mlir_module_string` into `mlir_module` with
// context `mlir_context`.
Status DeserializeMlirModule(llvm::StringRef serialized_mlir_module,
                             mlir::MLIRContext* mlir_context,
                             mlir::OwningOpRef<mlir::ModuleOp>* mlir_module);

}  // namespace tensorflow

  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;
  }
  return mlir_module;
}

}  // namespace test
}  // namespace mhlo

  llvm::raw_string_ostream os(serialized_mlir_module);
  mlir::OpPrintingFlags print_flags;
  print_flags.enableDebugInfo();
  module_op.print(os, print_flags);
  return std::move(os.str());
}

Status DeserializeMlirModule(llvm::StringRef serialized_mlir_module,
                             mlir::MLIRContext* mlir_context,
                             mlir::OwningOpRef<mlir::ModuleOp>* mlir_module) {

  StringRef host_module = op.getHostMlirModule();

  if (host_module.empty()) return success();

  mlir::OwningOpRef<mlir::ModuleOp> module_for_func;
  tensorflow::Status status = tensorflow::DeserializeMlirModule(
      host_module.str(), op->getContext(), &module_for_func);
  if (!status.ok()) {
    return op.emitError()
           << "attribute 'host_mlir_module' can not be deserialized. "
           << status.message();

  mlir::DialectRegistry mlir_registry;
  mlir::RegisterAllTensorFlowDialects(mlir_registry);
  mlir::MLIRContext mlir_context(mlir_registry);
  mlir::OwningOpRef<mlir::ModuleOp> mlir_module_with_add;
  TF_ASSERT_OK(DeserializeMlirModule(module_with_add, &mlir_context,
                                     &mlir_module_with_add));

  std::unique_ptr<mlir::Pass> partitioning_pass =
      mlir::TFTPU::CreateTPUResourceReadsWritesPartitioningPass();

  mlir::DialectRegistry registry;
  mlir::RegisterCommonToolingDialects(registry);
  mlir::MLIRContext context(registry);
  mlir::OwningOpRef<mlir::ModuleOp> mlir_module;

  TF_RETURN_IF_ERROR(
      DeserializeMlirModule(mlir_module_str, &context, &mlir_module));
  TF_RETURN_IF_ERROR(SetupReturnValues(*mlir_module, metadata_proto));
  TF_RETURN_IF_ERROR(SetupArguments(*mlir_module, arg_shapes, metadata_proto));

  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 mlir_module;
}

bool BuildAndRunPipeline(absl::string_view module_string,

    DialectRegistry mlir_registry;
    RegisterCommonToolingDialects(mlir_registry);
    context_.appendDialectRegistry(mlir_registry);
    TF_ASSERT_OK(
        tensorflow::DeserializeMlirModule(module_string, &context_, &module_));

    pm_ = std::make_unique<mlir::PassManager>(&context_);
    pm_->addNestedPass<mlir::func::FuncOp>(
        quant::stablehlo::CreateConvertTFQuantTypesPass());
  }

  mlir::DialectRegistry mlir_registry;
  RegisterAllTensorFlowDialects(mlir_registry);
  context->appendDialectRegistry(mlir_registry);
  mlir::OwningOpRef<mlir::ModuleOp> mlir_module;
  auto status =
      tensorflow::DeserializeMlirModule(string, context, &mlir_module);
  if (!status.ok()) {
    return status;
  }
  return mlir_module;
}

TEST(XlaRewriteUtilTest, TestEraseClusterFuncs) {
  static const char* const module_str =
      R"(