jax-ml/jax · error · ValueError

{dtype=} is not a valid dtype for JAX type promotion.

Error message

{dtype=} is not a valid dtype for JAX type promotion.

What it means

JAX's type-promotion lattice only accepts a fixed set of dtypes (its node set). When computing the least upper bound of a set of dtypes, one of the inputs is not a node in the lattice — typically because it is a NumPy dtype JAX does not register, such as a non-canonical extended or unsupported dtype.

Source

Thrown at jax/_src/dtypes.py:830

  # Note a potential algorithmic shortcut: from the definition of CUB(N), we have
  #   ∀ c ∈ N: CUB(N) ⊆ UB(c)
  # So if N ∩ CUB(N) is nonempty, if follows that LUB(N) = N ∩ CUB(N).
  N = set(nodes)
  if jax_numpy_dtype_promotion == config.NumpyDtypePromotion.STRICT:
    UB = _strict_lattice_ubs
  elif jax_numpy_dtype_promotion == config.NumpyDtypePromotion.STANDARD:
    if x64:
      UB = _standard_x64_lattice_ubs
    else:
      UB = _standard_x32_lattice_ubs
  else:
    raise ValueError(
      f"Unexpected value of jax_numpy_dtype_promotion={jax_numpy_dtype_promotion!r}")
  try:
    bounds = [UB[n] for n in N]
  except KeyError:
    dtype = next(n for n in N if n not in UB)
    raise ValueError(f"{dtype=} is not a valid dtype for JAX type promotion.")
  CUB = set.intersection(*bounds)
  LUB = (CUB & N) or {c for c in CUB if CUB.issubset(UB[c])}
  if len(LUB) == 1:
    return LUB.pop()
  elif len(LUB) == 0:
    if config.numpy_dtype_promotion.value == config.NumpyDtypePromotion.STRICT:
      msg = (
        f"Input dtypes {tuple(str(n) for n in nodes)} have no available implicit dtype "
        "promotion path when jax_numpy_dtype_promotion=strict. Try explicitly casting "
        "inputs to the desired output type, or set jax_numpy_dtype_promotion=standard.")
    elif any(n in _float8_dtypes for n in nodes):
      msg = (
        f"Input dtypes {tuple(str(n) for n in nodes)} have no available implicit dtype "
        "promotion path. To avoid unintended promotion, 8-bit floats do not support "
        "implicit promotion. If you'd like your inputs to be promoted to another type, "
        "you can do so explicitly using e.g. x.astype('float32')")
    elif any(n in _float6_dtypes for n in nodes):
      msg = (

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

  1. Inspect the dtypes of all inputs (x.dtype) and cast the offending one to a supported JAX dtype: x.astype(jnp.float32)
  2. Check that the value is not a non-numeric dtype (str/datetime/object) that leaked into an array
  3. Report/upgrade if you are using a supported dtype combination — a missing node may be a JAX bug

Example fix

# before
out = jnp.add(x_f16, y_custom_float8)

# after
out = jnp.add(x_f16.astype(jnp.float32), y_custom_float8.astype(jnp.float32))
Defensive patterns

Strategy: validation

Validate before calling

import jax.numpy as jnp, numpy as np
SUPPORTED = {np.dtype(d) for d in ['bool','uint8','uint16','uint32','uint64','int8','int16','int32','int64','float16','float32','float64','bfloat16','complex64','complex128']}
assert all(np.dtype(a.dtype) in SUPPORTED for a in arrays), 'unsupported dtype in promotion'

Try / catch

try:
    result_type = jnp.promote_types(a.dtype, b.dtype)
except ValueError:
    result_type = jnp.float32
    a, b = a.astype(result_type), b.astype(result_type)

Prevention

When it happens

Trigger: Calling jnp.promote_types or any binary op (e.g. jnp.add) where one operand's dtype is not in the promotion lattice, e.g. np.dtype('float16') mixed with an 8-bit custom float, or a bfloat16/extended dtype combination that was never registered, or dtypes like np.float128 / numpy string dtypes reaching promotion code.

Common situations: Mixing exotic NumPy dtypes (float128, datetime64, str) into JAX arrays; using custom/extended dtypes in ops that go through lattice_result_type; version upgrades that changed the set of lattice nodes.

Related errors


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