warp.dtype_to_jax

warp.dtype_to_jax(warp_dtype)

Return the Jax dtype corresponding to a Warp dtype.

Parameters:

warp_dtype – A Warp data type that has a corresponding Jax data type.

Raises:

TypeError – Unable to find a corresponding Jax data type.