warp.device_to_jax

warp.device_to_jax(warp_device)

Return the Jax device corresponding to a Warp device.

Only Jax devices that belong to the current process are returned. A Warp CUDA device maps to the local Jax GPU with the same process-local ordinal (jax.Device.local_hardware_id), which also holds in multi-process programs that use jax.distributed.initialize().

Returns:

jax.Device

Raises:

RuntimeError – Failed to find the corresponding Jax device in the current process.

Parameters:

warp_device (Device | str | None)