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)
addressable_memories(self)
default_memory(self)
get_stream_for_external_ready_events(self)
live_buffers(self)
memory(self, kind)
memory_stats(self)	Returns memory statistics for this device keyed by name.
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.
local_hardware_id	Opaque hardware ID, e.g., the CUDA device number.
platform
process_index	Integer index of this device's process.
task_id	Deprecated; please use process_index