warp.device_to_jax

warp.device_to_jax(warp_device)

Return the Jax device corresponding to a Warp device.

Returns:

jax.Device

Raises:

RuntimeError – Failed to find the corresponding Jax device.

Parameters:

warp_device (Device | str | None)