brainevent.jaxinfo_to_warpinfo

brainevent.jaxinfo_to_warpinfo(jax_info)

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

Takes a JAX shape-and-dtype specification and creates the corresponding Warp array type with matching data type and dimensionality. This is useful when building Warp kernel signatures from JAX output specifications.

Parameters:

jax_info (ShapeDtypeStruct) – A JAX structure containing shape, dtype, and ndim attributes describing an array.

Returns:

A Warp array type with matching data type and number of dimensions, suitable for use in Warp kernel definitions.

Return type:

Any

Raises:

• ImportError – If the warp package is not installed (propagated from jaxtype_to_warptype()).

• ValueError – If the dtype in jax_info is not supported by Warp (propagated from jaxtype_to_warptype()), or if jax_info describes a 0-D scalar (ndim == 0). Warp arrays have a minimum dimensionality of 1, so a 0-D input is rejected rather than silently promoted to ndim=1 (which would create a hidden JAX<->Warp shape mismatch).

See also

jaxtype_to_warptype

Convert a single dtype to a Warp type.

check_warp_installed

Verify that Warp is available.

Notes

The resulting Warp array type is constructed via warp.array(dtype=..., ndim=...) which creates a Warp type descriptor (not an actual array). This is typically used in Warp kernel function signatures to define input/output types.

Warp has a minimum array dimensionality of 1. A JAX 0-D scalar (shape == (), ndim == 0) therefore has no faithful Warp array representation; passing one raises ValueError instead of being silently bumped to ndim=1. Reshape the scalar to (1,) before conversion if a length-1 array is intended.

Examples

>>> import jax
>>> import numpy as np
>>> info = jax.ShapeDtypeStruct(shape=(32, 32), dtype=np.float32)
>>> warp_arr_type = jaxinfo_to_warpinfo(info)