Есть ли способ создать xla_client с процессором и графическим процессором?

Question

Я использую этот код в качестве основы. Но в этом xla_client я могу использовать только ресурсы графического процессора. Однако я хочу использовать память подкачки ЦП (унифицированную) в качестве ресурсов графического процессора. Как я могу создать xla_client с версией поддержки памяти ЦП? Там в любом случае..? Я действительно хочу создать клиент графического процессора с доступной памятью процессора.

      self.distributed_client = (
            xla_client._xla.get_distributed_runtime_client(
            server_address, host_id, use_coordination_service=False))
logger.debug(
        f"{host_id}: Trying to connect to xla runtime at {server_address}")
self.distributed_client.connect()
logger.debug(
        f"{host_id}: Success to connect to xla runtime at {server_address}")
if global_config.backend == "gpu":
    self.backend = xla_client.make_gpu_client(self.distributed_client,
                                                  node_id=host_id) 
else:
    raise NotImplementedError(
            f"backend {global_config.backend} is not supported")
self.local_devices = self.backend.local_devices()