jax-ml/jax · error · ValueError

Invalid argument to dtype: {x}.

Error message

Invalid argument to dtype: {x}.

What it means

jax.dtypes.dtype(x) (used broadly by jnp.dtype, default_int_dtype, supports_inf, etc.) rejects None as an argument because there is no meaningful dtype. The f-string renders the None value.

Source

Thrown at jax/_src/dtypes.py:1005

def register_type_whose_dtype_should_not_be_canonicalized(typ: type):
  global _types_whose_dtype_should_not_be_canonicalized
  _types_whose_dtype_should_not_be_canonicalized += (typ,)

def dtype(x: Any) -> DType:
  """Return the dtype object for a value or type.

  Python scalars, Python scalar types, NumPy scalar type, NumPy dtypes, and
  non-JAX arrays will have their dtypes canonicalized.

  Note: this is not the same function as jax.numpy.dtype, which simply aliases
  numpy.dtype."""
  # TODO(phawkins): in the future, we would like to:
  # - return the default dtype for Python scalar types and values
  # - canonicalize NumPy array and scalar types
  # - return NumPy dtypes as-is, uncanonicalized.
  if x is None:
    raise ValueError(f"Invalid argument to dtype: {x}.")
  if isinstance(x, type):
    # Python scalar types, e.g., int, float
    if (dt := python_scalar_types_to_dtypes.get(x)) is not None:
      return canonicalize_dtype(dt)

    # 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

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Guard before calling: dt = dtype(x) if x is not None else jnp.float32
2. Use jnp.result_type which substitutes default_float_dtype for None
3. Pass a concrete dtype or omit the argument instead of None

Example fix

# before
dt = jax.dtypes.dtype(maybe_none_dtype)

# after
dt = jax.dtypes.dtype(maybe_none_dtype) if maybe_none_dtype is not None else jnp.result_type(1.0)

if x is None:
    x = 1.0  # or raise your own descriptive error
dt = jax.dtypes.dtype(x)

def is_inferable_dtype_arg(x) -> bool:
    return x is not None

When it happens

Trigger: Passing None directly: jax.dtypes.dtype(None), jnp.zeros(3, dtype=None) paths where None is not treated as 'use default', or a helper that forwards an unset dtype variable into dtype().

Common situations: A config object with an optional dtype field (None default) forwarded without a fallback; refactoring where dtype=None previously meant default in old JAX versions but now must be omitted.

Understand the failure class

Background: Invalid argument type errors: "must be of type string", "expected X, got Y", and ERR_INVALID_ARG_TYPE explained — this error's family across 15 libraries.

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/53e31b9a00f35f85. Report an issue: GitHub.