jax-ml/jax · error · TypeError

JAX only supports number, bool, and string dtypes, got dtype

Error message

JAX only supports number, bool, and string dtypes, got dtype {dtype} in {fun_name}

What it means

check_and_canonicalize_user_dtype resolved the requested dtype to a NumPy dtype whose kind is not number/bool/string (e.g. datetime, timedelta, object, void). Unless allow_non_jax_dtypes was enabled, JAX rejects it, naming the operation via fun_name.

Source

Thrown at jax/_src/dtypes.py:1129

  if dtype is None:
    raise ValueError("dtype must be specified.")
  if isinstance(dtype, Array):
    raise ValueError("Passing an array as a dtype argument is no longer "
                     "supported; instead of dtype=arr use dtype=arr.dtype.")
  if issubdtype(dtype, extended):
    return dtype
  # Avoid using `dtype in [...]` because of numpy dtype equality overloading.
  if isinstance(dtype, type) and (f := _DEFAULT_TYPEMAP.get(dtype)) is not None:
    return f()
  np_dtype = np.dtype(dtype)
  if np_dtype not in _jax_dtype_set:
    if allow_non_jax_dtypes:
      return np_dtype
    msg = (
        f'JAX only supports number, bool, and string dtypes, got dtype {dtype}'
    )
    msg += f" in {fun_name}" if fun_name else ""
    raise TypeError(msg)
  return _maybe_canonicalize_explicit_dtype(np_dtype, fun_name or "")

def safe_to_cast(input_dtype_or_value: Any,
                 output_dtype_or_value: Any) -> bool:
  """Check if a dtype/value is safe to cast to another dtype/value

  Args:
    input_dtype_or_value: a dtype or value (to be passed to result_type)
      representing the source dtype.
    output_dtype_or_value: a dtype or value (to be passed to result_type)
      representing the target dtype.

  Returns:
    boolean representing whether the values are safe to cast according to
    default type promotion semantics.

  Raises:
    TypePromotionError: if the inputs have differing types and no type promotion

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Convert datetimes to numeric: arr.astype('int64') or divide to float seconds
2. Use numeric/bool/string dtypes only: jnp.float32, jnp.int32, jnp.bool_
3. Strip structured arrays to their fields: arr['field'].astype('float32')

Example fix

# before
lax.conv_general_dilated(..., dtype=np.dtype('datetime64[ns]'))

# after
lax.conv_general_dilated(..., dtype=jnp.float32)  # after converting data

import numpy as np
np_dt = np.dtype(dtype)
assert np_dt.kind in 'biu b'.replace(' ',''), f'dtype {np_dt} kind {np_dt.kind!r} not allowed'  # numeric/bool
assert np_dt.kind in 'biu', 'use numeric/bool/string dtypes only'

def is_allowed_user_dtype(dt) -> bool:
    import numpy as np
    try:
        return np.dtype(dt).kind in 'biuUS'
    except TypeError:
        return False

When it happens

Trigger: dtype=np.dtype('datetime64[ns]') or 'O' passed to ops like conv_general_dilated / searchsorted / sds wrappers; structured dtypes (kind 'V'); timedelta64 inputs from time pipelines.

Common situations: Timestamp features not converted to numeric before model code; structured arrays from file I/O (HDF5/npz with compound dtypes) fed into lax ops; passing dtype strings from external schemas.

Related errors

• Dtype {dtype} is not a valid JAX array type. Only arrays of
• dtype cannot be None.
• unexpected input: {dtype=}
• Unrecognized {kind=} expected one of {list(_dtype_kinds.keys
• Value '{x}' with dtype {dt} is not a valid JAX array type. O

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/c2f9506eaa459055. Report an issue: GitHub.