std::string GetPjRtDeviceCompilationProfilerResourceName(
    const DeviceType& device_type);

// Gets the ResourceMgr where the DeviceCompiler is/should be stored for the
// given `device_type`.
absl::StatusOr<ResourceMgr*> GetResourceMgrForDeviceCompiler(
    const OpKernelContext& ctx, const DeviceType& device_type);

}  // namespace tensorflow

    bool may_alias_resource_update,
    const XlaCompiler::CompilationResult** compilation_result,
    xla::PjRtClient** client, xla::PjRtLoadedExecutable** executable) {
  TF_ASSIGN_OR_RETURN(ResourceMgr * rm, GetResourceMgrForDeviceCompiler(
                                            ctx, platform_info.device_type()));
  return CompileToPjRtLoadedExecutable(
      ctx.device(), platform_info, function, args, compile_mode, has_ref_vars,

    const DeviceType& device_type) {
  return absl::StrCat(kPjRtDeviceCompilationProfilerResourceName, "_",
                      device_type.type_string());
}

absl::StatusOr<ResourceMgr*> GetResourceMgrForDeviceCompiler(
    const OpKernelContext& ctx, const DeviceType& device_type) {
  // We store information about the JIT-compiled XLA computation in the
  // ResourceMgr. The DeviceCompiler (which contains the DeviceCompilationCache)

    FunctionLibraryRuntime* flr,
    DeviceCompiler<xla::PjRtLoadedExecutable, xla::PjRtClient>**
        pjrt_device_compiler,
    DeviceCompilationProfiler** profiler) {
  TF_ASSIGN_OR_RETURN(ResourceMgr * rm, GetResourceMgrForDeviceCompiler(
                                            ctx, platform_info.device_type()));
  return GetOrCreatePjRtDeviceCompilerAndProfiler(
      platform_info, rm, flr, pjrt_device_compiler, profiler);
}