Device

Auto-generated from jax_code_defs

PyTorchJAX (Core)Apple MLX

PyTorch

API: torch.device

Strategy: Plugin (device_allocator)

JAX (Core)

API: jax.devices

Strategy: Plugin (device_allocator)

Apple MLX

API: mlx.core.Stream.device

Strategy: Direct Mapping

Official Docs ↗