DTypeLike

brainstate.typing.DTypeLike

Union of types that can be converted to a valid JAX dtype.

This is more restrictive than numpy.typing.DTypeLike as JAX doesn’t support object arrays or structured dtypes. It excludes None to require explicit handling of optional dtypes.

Components

str

String representations like ‘float32’, ‘int32’, ‘bool’.

type[Any]

Type objects like np.float32, float, int, bool.

np.dtype

NumPy dtype objects created with np.dtype().

SupportsDType

Any object with a .dtype property.

Examples

>>> def cast_array(array: ArrayLike, dtype: DTypeLike) -> jax.Array:
...     '''Cast array to specified dtype.'''
...     return jnp.asarray(array, dtype=dtype)
>>>
>>> # Valid dtype specifications
>>> cast_array(data, 'float32')           # String
>>> cast_array(data, np.float32)          # NumPy type
>>> cast_array(data, float)               # Python type
>>> cast_array(data, np.dtype('int32'))   # NumPy dtype object
>>> cast_array(data, other_array)         # Object with dtype property