params.device = device;
  OpKernelContext ctx(&params, 0);

  PjRtDeviceCompiler* pjrt_device_compiler = nullptr;
  DeviceCompilationProfiler* profiler = nullptr;
  TF_EXPECT_OK(GetOrCreatePjRtDeviceCompilerAndProfiler(
      ctx, platform_info, device_setup_.flr(), &pjrt_device_compiler,
      &profiler));
  core::ScopedUnref pjrt_device_compiler_ref(pjrt_device_compiler);
  core::ScopedUnref profiler_ref(profiler);

  *xla_device_compiler = CreateXlaDeviceCompiler(
      persistor_config, compilation_device_type, client);
  return absl::OkStatus();
}

Status GetOrCreatePjRtDeviceCompilerAndProfiler(
    const XlaPlatformInfo& platform_info, ResourceMgr* rm,
    FunctionLibraryRuntime* flr, PjRtDeviceCompiler** pjrt_device_compiler,
    DeviceCompilationProfiler** profiler) {

    xla::PjRtClient** client, xla::PjRtLoadedExecutable** executable) {
  PjRtDeviceCompiler* pjrt_device_compiler;
  DeviceCompilationProfiler* profiler;
  TF_RETURN_IF_ERROR(GetOrCreatePjRtDeviceCompilerAndProfiler(
      platform_info, rm, flr, &pjrt_device_compiler, &profiler));

  // Hold the reference to the PJRT device compiler and profiler during

// non-XLA devices aren't supported yet. This is because:
// 1. PjRtClient doesn't support data transfer for non-XLA devices yet
// 2. Fetching the PjRtClient for non-XLA devices is also not supported yet
Status GetOrCreatePjRtDeviceCompilerAndProfiler(
    const OpKernelContext& ctx, const XlaPlatformInfo& platform_info,
    FunctionLibraryRuntime* flr,
    DeviceCompiler<xla::PjRtLoadedExecutable, xla::PjRtClient>**
        pjrt_device_compiler,

    PjRtDeviceCompiler** pjrt_device_compiler,
    DeviceCompilationProfiler** profiler,
    const XlaCompiler::CompilationResult** result,
    xla::PjRtLoadedExecutable** executable) {
  TF_RETURN_IF_ERROR(GetOrCreatePjRtDeviceCompilerAndProfiler(
      *ctx, platform_info_, ctx->function_library(), pjrt_device_compiler,
      profiler));

  XlaCompiler::Options options =