DeviceType device_cpu = DeviceType(DEVICE_CPU);
  EXPECT_EQ(GetPersistentCacheDirectory(device_cpu), "/tmp/xla_cache");
  DeviceType device_tpu = DeviceType(DEVICE_TPU);
  EXPECT_TRUE(GetPersistentCacheDirectory(device_tpu).empty());
}

TEST_F(XlaPlatformInfoTest, GetPersistentCacheDirectoryNoDeviceTypes) {
  tensorflow::GetMarkForCompilationPassFlags()

    TF_ASSIGN_OR_RETURN(*pjrt_client, GetOrCreatePjRtClient(device_type));
    return absl::OkStatus();
  }

  // TFRT-TPU is used if device_type is `DEVICE_TPU` and platform_info does not
  // have `xla_device_metadata`.
  if (device_type == DEVICE_TPU) {
    *compilation_device_type = DeviceType(DEVICE_TPU_XLA_JIT);
    TF_ASSIGN_OR_RETURN(*pjrt_client, GetOrCreatePjRtClient(device_type));
    return absl::OkStatus();

  EXPECT_FALSE(UsePjRtForSingleDeviceCompilation(DeviceType(DEVICE_CPU)));
}

TEST(XlaCompileUtilTest, PjRtDeviceCompilerResourceName) {
  EXPECT_EQ(GetPjRtDeviceCompilerResourceName(DeviceType(DEVICE_TPU)),
            "pjrt_device_compiler_TPU");
  EXPECT_EQ(GetPjRtDeviceCompilerResourceName(DeviceType(DEVICE_TPU_NODE)),
            "pjrt_device_compiler_TPU");

                             compiler_arg_source));

  XlaPlatformInfo platform_info = XlaPlatformInfoFromDevice(dev);
  auto compilation_device_type = platform_info.device_type();
  if (platform_info.device_type() != DEVICE_TPU) {
    TF_ASSIGN_OR_RETURN(compilation_device_type,
                        GetCompilationDeviceType(platform_info.device_type()));
  }

  XlaDeviceCompiler* xla_device_compiler;

  // for CPU/GPU. This is to make sure the DeviceCompiler's lifecycle is
  // maintained appropriately.
  ResourceMgr* rm = nullptr;
  if (device_type == DEVICE_TPU) {
    rm = tfrt_global::GetTFGlobalResourceMgr();
  } else {
    rm = ctx.resource_manager();
  }

  if (!rm) {
    return absl::InternalError("No resource manager found.");
  }
  return rm;
}

  std::vector<xla::XlaOp> returns(1);
  return BuildHloFromGraph(graph, builder, mlir_context, xla_params, returns,
                           use_output_shapes, /*args=*/{},
                           /*control_rets=*/{}, DEVICE_TPU,
                           FunctionLibraryDefinition(OpRegistry::Global()),
                           /*debug_info=*/{},
                           /*custom_legalization_passes=*/{});
}