protected:
  void SetUp() override {
    xla_compiler_client_ = std::make_unique<XlaDeviceCompilerClient>(
        xla::ClientLibrary::LocalClientOrDie());
    TF_ASSERT_OK(CreatePjRtCompilerClient());

    XlaOpRegistry::RegisterCompilationKernels();

    flib_def_ = std::make_unique<FunctionLibraryDefinition>(
        OpRegistry::Global(), FunctionDefLibrary());

    cache_dir_ = testing::TmpDir();