brainevent.jaxtype_to_warptype

brainevent.jaxtype_to_warptype(dtype)

Convert a JAX / NumPy dtype to the corresponding Warp scalar type.

Maps standard NumPy data types (which are also used by JAX) to their Warp equivalents. This is needed when constructing Warp kernel signatures or Warp array types from JAX metadata.

Parameters:

dtype (dtype | type) – The data type to convert. Accepts anything numpy.dtype can normalize, including NumPy scalar types (np.float32), JAX dtypes (jnp.float32), np.dtype instances (np.dtype('float32')), and Python builtins (float, int, bool). Builtins are normalized via numpy.dtype first (e.g. float -> float64, int -> the platform default integer, bool -> bool_).

Returns:

The corresponding Warp scalar type (e.g., warp.float32, warp.int32, warp.bool).

Return type:

Any

Raises:

• ImportError – If the warp package is not installed.

• ValueError – If dtype cannot be normalized by numpy.dtype or does not correspond to any supported Warp type. Supported types are: float16, float32, float64, int8, int16, int32, int64, uint8, uint16, uint32, uint64, and bool_. bfloat16 and complex dtypes are not supported by Warp and raise a clear ValueError.

See also

jaxinfo_to_warpinfo

Convert a full jax.ShapeDtypeStruct to a Warp array type.

check_warp_installed

Verify that Warp is available.

Notes

The dtype is first normalized with numpy.dtype(dtype) so that Python builtins and dtype-like objects all resolve to a canonical numpy.dtype before mapping. This avoids the pitfall that float == np.float64 is False even though np.dtype(float) is float64. The mapping covers all scalar types supported by both NumPy and Warp. bfloat16 (which Warp has no scalar for) and complex types are unsupported and raise ValueError.

Examples

>>> import numpy as np
>>> warp_type = jaxtype_to_warptype(np.float32)
>>> warp_type  # warp.float32
>>> jaxtype_to_warptype(float) is jaxtype_to_warptype(np.float64)
True