jax-ml/jax · error · ValueError

Passing an array as a dtype argument is no longer supported;

Error message

Passing an array as a dtype argument is no longer supported; instead of dtype=arr use dtype=arr.dtype.

What it means

Older JAX versions accepted a 0-d array as the dtype= argument (using its value as a dtype code); this was removed. Passing a jax.Array (e.g. jnp.int32(0) or np.array(3)) as dtype now raises immediately with a hint to use arr.dtype instead.

Source

Thrown at jax/_src/dtypes.py:1114

  if len(args) == 0:
    raise ValueError("at least one array or dtype is required")
  dtype: DType | ExtendedDType
  dtype, weak_type = lattice_result_type(*(default_float_dtype() if arg is None else arg for arg in args))
  if weak_type:
    dtype = default_types['f' if dtype in _custom_float_dtypes else dtype.kind]()
  return (dtype, weak_type) if return_weak_type_flag else dtype

def check_and_canonicalize_user_dtype(
    dtype, fun_name=None, *, allow_non_jax_dtypes: bool = False
) -> DType:
  """Checks validity of a user-provided dtype, and returns its canonical form.

  For Python scalar types this function returns the corresponding default dtype.
  """
  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,

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Use the dtype class: dtype=jnp.int32 / np.float32
2. Use the array's dtype attribute: dtype=arr.dtype
3. If selecting from a registry, store dtype objects/classes, not arrays

Example fix

# before
op(..., dtype=jnp.array(7))  # old-style numeric dtype code

# after
op(..., dtype=jnp.int32)

from jax import Array
if isinstance(dtype, Array):
    dtype = dtype.dtype
op(..., dtype=dtype)

import numpy as np
from jax import Array

def is_valid_dtype_arg(dt: object) -> bool:
    return not isinstance(dt, Array) and not isinstance(dt, np.ndarray)

When it happens

Trigger: dtype=jnp.array(3), dtype=np.array('float32'), or forwarding a computed scalar array into a dtype= kwarg of ops routed through check_and_canonicalize_user_dtype (conv_general_dilated, searchsorted, sds functions, etc.).

Common situations: Legacy code or tutorials written for old JAX; dynamically-built dtype arguments stored as arrays; copy-paste from NumPy where np.dtype(np.array(...)) sometimes worked.

Related errors

• primal and tangent arguments to jax.jvp do not match; dtypes
• reduce_axes argument to vjp is deprecated
• 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

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