Device

Constructs a device object.

Abstract Signature:

Device(type: str, index: int)

PyTorchJAX (Core)NumPyKerasTensorFlowApple MLXFlax NNXPaxML / Praxis

PyTorch

API: torch.nn.parallel.distributed.DistributedDataParallel.device

Strategy: Plugin (device_allocator)

JAX (Core)

API: —

Strategy: Plugin (device_allocator)

NumPy

API: —

Strategy: Macro ''cpu''

Keras

API: —

Strategy: Plugin (device_allocator)

TensorFlow

API: —

Strategy: Plugin (device_allocator)

Apple MLX

API: mlx.core.Stream.device

Strategy: Plugin (device_allocator)

Flax NNX

API: —

Strategy: Plugin (device_allocator)

PaxML / Praxis

API: —

Strategy: Plugin (device_allocator)