}

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

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

  Status LegalizeSingleOp(Operation& op) {
    SourceMgrDiagnosticHandler sourceMgrHandler(source_manager_, &context_);

    OpBuilder op_builder(&op);

  registry.insert<mlir::func::FuncDialect, mlir::tf_type::TFTypeDialect,
                  mlir::quant::QuantizationDialect>();
  mlir::quant::RegisterOps();

  MLIRContext context(registry);
  context.loadAllAvailableDialects();

  llvm::SourceMgr source_mgr;
  SourceMgrDiagnosticHandler sourceMgrHandler(source_mgr, &context);

  auto verbose_dir = llvm::sys::path::parent_path(output_path).str();

namespace {

void RegisterDialects(mlir::MLIRContext& ctx) {
  mlir::DialectRegistry registry;
  mlir::RegisterAllTensorFlowDialects(registry);
  ctx.appendDialectRegistry(registry);
  ctx.loadAllAvailableDialects();
}

Status ConvertDataTypeToTensor(tensorflow::DataType dtype, Builder builder,
                               Type* type) {
  Status s = tensorflow::ConvertDataType(dtype, builder, type);

    dialects.insert<TF::TensorFlowDialect, func::FuncDialect, chlo::ChloDialect,
                    mhlo::MhloDialect, quant::QuantizationDialect>();
    ctx_ = std::make_unique<MLIRContext>(dialects);
    ctx_->loadAllAvailableDialects();

    // Create a CPU client with 1 device.
    TF_ASSERT_OK_AND_ASSIGN(
        pjrt_client_,
        xla::GetTfrtCpuClient(/*asynchronous=*/false, /*cpu_device_count=*/1));