jax-ml/jax · error · TypeError

Value '{x}' with dtype {dt} is not a valid JAX array type. O

Error message

Value '{x}' with dtype {dt} is not a valid JAX array type. Only arrays of numeric types are supported by JAX.

What it means

When dtype() is given a string or np.dtype, it converts via np.dtype(x) and verifies the result is a JAX-supported type (numeric/bool or an extended dtype). Strings that parse to non-numeric NumPy dtypes — like 'str', 'U10', 'datetime64', 'object' — fail this check.

Source

Thrown at jax/_src/dtypes.py:1028

    # Numpy scalar types, e.g., np.int32, np.float32
    if _issubclass(x, np.generic):
      dt = np.dtype(x)
      return _maybe_canonicalize_explicit_dtype(dt, "dtype")

  # Python scalar values, e.g., int(3), float(3.14)
  elif (dt := python_scalar_types_to_dtypes.get(type(x))) is not None:
    return canonicalize_dtype(dt)
  # Jax Arrays, literal arrays, and scalars.
  # We intentionally do not canonicalize these types: once we've formed an x64
  # value, that is something we respect irrespective of the x64 mode.
  elif isinstance(x, _types_whose_dtype_should_not_be_canonicalized):
    return x.dtype

  if isinstance(x, (str, np.dtype)):
    dt = np.dtype(x)
    if dt not in _jax_dtype_set and not issubdtype(dt, extended):
      raise TypeError(f"Value '{x}' with dtype {dt} is not a valid JAX array "
                      "type. Only arrays of numeric types are supported by JAX.")
    return _maybe_canonicalize_explicit_dtype(dt, "dtype")

  # If x has a dtype attribute, and it's a valid dtype, use it. This avoids
  # calling np.result_type on objects that might have a .dtype but are not
  # standard NumPy array-like, which can lead to warnings in NumPy 2.4+.
  dt_attr = getattr(x, 'dtype', None)
  if issubdtype(dt_attr, extended) or isinstance(dt_attr, np.dtype):
    dt = dt_attr
  else:
    try:
      dt = np.result_type(x)
    except TypeError as err:
      raise TypeError(f"Cannot determine dtype of {x}") from err
  if dt not in _jax_dtype_set and not issubdtype(dt, extended):
    raise TypeError(f"Value '{x}' with dtype {dt} is not a valid JAX array "
                    "type. Only arrays of numeric types are supported by JAX.")
  # TODO(jakevdp): fix return type annotation and remove this ignore.

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Map string data to numeric ids first, then use an integer dtype
Use a supported dtype string: 'float32','int32','complex64','bool', etc.
If you genuinely need extended dtypes (e.g. jax.dtypes.prng_key), use the dtype object, not a string

Example fix

# before
jax.dtypes.dtype('U10')

# after
ids = np.array([s.encode() for s in strings], dtype=np.int32)  # encode first
jax.dtypes.dtype(ids.dtype)

Defensive patterns

Strategy: type-guard

Validate before calling

import numpy as np
if isinstance(x, (str, np.dtype)):
    assert np.dtype(x).kind in 'biufc', f'non-numeric dtype {np.dtype(x)}'

Type guard

def is_numeric_dtype_spec(s: object) -> bool:
    import numpy as np
    if not isinstance(s, (str, np.dtype)):
        return True
    try:
        return np.dtype(s).kind in 'biufc'
    except TypeError:
        return False

Prevention

Encode strings to integer ids before entering JAX code
Validate dtype strings against a whitelist before passing them to dtype()

When it happens

Trigger: jax.dtypes.dtype('U10'), dtype('S'), dtype('datetime64[ns]'), dtype('object'), or jnp.array(..., dtype='str') style paths that reach this validation.

Common situations: Passing a format string meant for something else as a dtype; encoding text data; copying NumPy snippets that use 'datetime64' dtypes; column-type strings from a schema (e.g. 'object' from pandas).

Related errors

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