jax.devices#

jax.devices(backend=None)[source]#

Returns a list of all devices for a given backend.

Each device is represented by a subclass ofDevice (e.g.CpuDevice,GpuDevice). The length of the returned list isequal todevice_count(backend). Local devices can be identified bycomparingDevice.process_index to the value returned byjax.process_index().

Ifbackend isNone, returns all the devices from the default backend.The default backend is generally'gpu' or'tpu' if available,otherwise'cpu'.

Parameters:

backend (str |xla_client.Client |None) – This is an experimental feature and the API is likely to change.Optional, a string representing the xla backend:'cpu','gpu', or'tpu'.

Returns:

List of Device subclasses.

Return type:

list[xla_client.Device]