jax-ml/jax · error · TypeError

dtype {dtype!r} not understood

Error message

dtype {dtype!r} not understood

What it means

jax.dtypes.canonicalize_dtype could not interpret the value as a numpy dtype — np.dtype(dtype) itself raised TypeError, so the object is not dtype-like (e.g. a random class instance, list, or malformed string).

Source

Thrown at jax/_src/dtypes.py:358

def to_complex_dtype(dtype: DTypeLike) -> DType:
  ftype = to_inexact_dtype(dtype)
  if ftype in [np.dtype('float64'), np.dtype('complex128')]:
    return np.dtype('complex128')
  return np.dtype('complex64')


@functools.cache
def _canonicalize_dtype(x64_enabled: bool, allow_extended_dtype: bool, dtype: Any) -> DType | ExtendedDType:
  if issubdtype(dtype, extended):
    if not allow_extended_dtype:
      raise ValueError(f"Internal: canonicalize_dtype called on extended dtype {dtype} "
                       "with allow_extended_dtype=False")
    return dtype
  try:
    dtype_ = np.dtype(dtype)
  except TypeError as e:
    raise TypeError(f'dtype {dtype!r} not understood') from e

  if x64_enabled:
    return dtype_
  else:
    return _dtype_to_32bit_dtype.get(dtype_, dtype_)

@overload


def canonicalize_dtype(
    dtype: Any, allow_extended_dtype: Literal[False] = False
) -> DType:
  ...


@overload

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Pass standard dtype specifiers: 'float32', np.float32, jnp.bfloat16
2. Validate with np.dtype(dtype) in a try/except before handing user config to JAX
3. Fix the source producing non-dtype objects (config parsing, default values)

Example fix

# before
jit(fn)(x, dtype=config['dtype'])  # dtype is a dict

# after
dtype = np.dtype(config['dtype'])  # validate early
jit(fn)(x, dtype=dtype)

try:
    np.dtype(dtype)
except TypeError:
    raise ValueError(f'invalid dtype {dtype!r}') from None

import numpy as np
def is_dtype_like(v) -> bool:
    try:
        np.dtype(v); return True
    except TypeError:
        return False

try:
    canonicalize_dtype(dtype)
except TypeError:
    dtype = np.float32  # or reject input

When it happens

Trigger: canonicalize_dtype(SomeObject()), passing an invalid dtype string like 'float128' on platforms without it, or a dataclass where __repr__/conversion breaks np.dtype.

Common situations: User-supplied dtype fields from YAML/config parsed into non-dtype objects; typos in dtype names in kwargs like jit(f, ...)? or custom primitives.

Related errors

• primal and tangent arguments to jax.jvp do not match; dtypes
• unexpected JAX type (e.g. shape/dtype) for gradient ref pass
• Accumulator aval mismatch: expected {aval}, got {acc.aval}
• unexpected JAX type (e.g. shape/dtype) for argument to VJP f
• linear_transpose only supports [float or complex] -> [float

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