jax.Device

class jax.Device

可用设备的描述符。

子类用于表示特定类型的设备，例如 CPU、GPU。子类可能具有特定于该设备类型的其他属性。

__init__(*args, **kwargs)

方法

__init__(*args, **kwargs)

属性

addressable_memories	返回设备可以寻址的所有内存。
client	(self) -> object
default_memory	返回设备的默认内存。
device_kind	(self) -> str
get_stream_for_external_ready_events
host_id	已弃用；请使用 process_index
id	此设备的整数 ID。
live_buffers
local_hardware_id	不透明的硬件 ID，例如 CUDA 设备编号。
memory
memory_stats	返回此设备的内存统计信息，以名称为键。
platform	(self) -> str
process_index	此设备进程的整数索引。
task_id	已弃用；请使用 process_index
transfer_from_outfeed
transfer_to_infeed