jax.local_devices

jax.local_devices(process_index=None, backend=None, host_id=None)

與 jax.devices() 類似，但僅傳回給定進程的本機裝置。

如果 process_index 為 None，則傳回此進程的本機裝置。

參數:

• process_index (int | None) – 進程的整數索引。進程索引可以透過 len(jax.process_count()) 取得。

• backend (str | xla_client.Client | None) – 這是一項實驗性功能，API 可能會變更。選用，代表 xla 後端的字串：'cpu'、'gpu' 或 'tpu'。

• host_id (int | None)

傳回:

裝置子類別的清單。

傳回類型:

list[xla_client.Device]