warp.device_from_jax

warp.device_from_jax(jax_device)

Return the Warp device corresponding to a Jax device.

A Jax GPU maps to the Warp CUDA device with the same process-local ordinal (jax.Device.local_hardware_id). This differs from the global jax.Device.id in multi-process programs that use jax.distributed.initialize().

Parameters:

jax_device (jax.Device) – A Jax device descriptor.

Raises:

RuntimeError – The Jax device is neither a CPU nor GPU device, or it is a GPU that belongs to another process.

Return type:

Device