jax.devices

jax.devices(backend=None)

傳回給定後端的所有裝置列表。

每個裝置都由 Device (例如 CpuDevice、GpuDevice) 的子類別表示。傳回列表的長度等於 device_count(backend)。可以透過比較 Device.process_index 與 jax.process_index() 傳回的值來識別本機裝置。

如果 backend 為 None，則傳回預設後端的所有裝置。如果可用，預設後端通常為 'gpu' 或 'tpu'，否則為 'cpu'。

參數:

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

傳回:

裝置子類別的列表。

傳回類型:

list[xla_client.Device]