jax-ml/jax · error · TypeError

Dtype {dtype} is not a valid JAX array type. Only arrays of

Error message

Dtype {dtype} is not a valid JAX array type. Only arrays of numeric types are supported by JAX.

What it means

check_valid_dtype rejects dtypes that are not in JAX's supported set (_jax_dtype_set). JAX only supports numeric (and bool) array element types; this fires when converting a traced/numpy value whose dtype is e.g. object, string, datetime, or a NumPy type JAX does not handle.

Source

Thrown at jax/_src/dtypes.py:954


def is_weakly_typed(x: Any) -> bool:
  if type(x) in _weak_types or type(x) in _registered_weak_types:
    return True
  try:
    return x.aval.weak_type
  except AttributeError:
    return False

def is_weakly_typed_scalar(x: Any) -> bool:
  try:
    return x.aval.weak_type and np.ndim(x) == 0
  except AttributeError:
    return type(x) in python_scalar_types

def check_valid_dtype(dtype: DType) -> None:
  if dtype not in _jax_dtype_set:
    raise TypeError(f"Dtype {dtype} is not a valid JAX array "
                    "type. Only arrays of numeric types are supported by JAX.")

def _maybe_canonicalize_explicit_dtype(dtype: DType, fun_name: str) -> DType:
  "Canonicalizes explicitly requested dtypes, per explicit_x64_dtypes."
  allow = config.explicit_x64_dtypes.value
  if allow == config.ExplicitX64Mode.ALLOW or config.enable_x64.value:
    return dtype
  canonical_dtype = canonicalize_dtype(dtype)
  if canonical_dtype == dtype:
    return dtype
  fun_name = f" requested in {fun_name}" if fun_name else ""
  if allow == config.ExplicitX64Mode.ERROR:
    msg = ("Explicitly requested dtype {}{} is not available. To enable more "
           "dtypes, set the jax_enable_x64 or allow_explicit_x64_dtypes "
           "configuration options."
          "See https://github.com/jax-ml/jax#current-gotchas for more.")
    msg = msg.format(dtype, fun_name, canonical_dtype.name)
    raise ValueError(msg)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

  1. Convert the array to a numeric dtype before passing it: np.asarray(x, dtype=np.float32)
  2. Check x.dtype on the offending input and fix upstream data loading (e.g. pd.to_numeric)
  3. For string data, use a separate tokenizer/encoding step instead of JAX arrays

Example fix

# before
arr = np.array(['1', '2', '3'])
jnp.sin(arr)

# after
arr = np.array([1, 2, 3], dtype=np.float32)
jnp.sin(arr)
Defensive patterns

Strategy: type-guard

Validate before calling

import numpy as np
def is_jax_numeric(arr) -> bool:
    return getattr(arr, 'dtype', None) is not None and np.asarray(arr).dtype.kind in 'biufc'

Type guard

def has_valid_jax_dtype(x) -> bool:
    dt = getattr(x, 'dtype', None)
    return dt is not None and (str(dt) in {'bool','int8','int16','int32','int64','uint8','uint16','uint32','uint64','float16','float32','float64','complex64','complex128'} or dt.kind in 'biufc')

Prevention

When it happens

Trigger: Passing np.array(['a','b']) or an object-dtype array to a jnp function; arrays of np.datetime64; np.void/record dtypes; using a custom NumPy dtype subclass when creating ShapedArrays (via _make_shaped_array_for_numpy_array / numpy_scalar paths, e.g. inside jit tracing of constants).

Common situations: Accidentally loading a CSV column of strings into an array fed to JAX; object arrays from pandas; mixed-type lists that NumPy upcasts to object dtype.

Related errors

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