jax.Device#
- class jax.Device#
A descriptor of an available device.
Subclasses are used to represent specific types of devices, e.g. CPUs, GPUs. Subclasses may have additional properties specific to that device type.
- __init__(*args, **kwargs)#
Methods
__init__
(*args, **kwargs)live_buffers
(self)transfer_from_outfeed
(self, arg0)transfer_to_infeed
(self, arg0)Attributes
client
device_kind
host_id
Deprecated; please use process_index
id
Integer ID of this device.
platform
process_index
Integer index of this device's process.
task_id
Deprecated; please use process_index