jax-ml/jax · error · ValueError

Cannot convert_element_type from

Error message

Cannot convert_element_type from {dtype_to_string(x.dtype)} to {dtype_to_string(edtype)}

What it means

When converting an array to an extended dtype, the dtype's rules must allow the conversion (via a convert_to callback or an allow_conversion flag). If the source dtype is not an allowed input for that extended dtype (e.g. converting int32 directly to a float8 type that only accepts its representation dtype), ValueError is raised.

Solutions

Convert to the extended dtype's representation dtype first (e.g. jnp.float32), then to the extended dtype
Check edtype._rules.allow_conversion / convert_to to see which pairs are permitted
Use the documented supported source dtype for the target extended dtype
For custom dtypes, extend convert_to to accept the pair if semantically valid

Example fix

// before
x_f8 = x.astype(dtypes.float8_e4m3fn)  # from int32, disallowed

// after
x_f8 = x.astype(jnp.float32).astype(dtypes.float8_e4m3fn)

Defensive patterns

Strategy: validation

Validate before calling

def conversion_allowed(src_dtype, edtype):
    fn = getattr(edtype._rules, 'convert_to', None)
    return fn(src_dtype, edtype) if fn else edtype._rules.allow_conversion

Try / catch

try:
    x_ed = x.astype(edtype)
except ValueError:
    x_ed = x.astype(jnp.float32).astype(edtype)

Prevention

Route conversions through the documented intermediate dtype (usually float32)
Check dtype rules before astype in dtype-generic code
Wrap user-facing dtype arguments in an allowlist

When it happens

Trigger: jax.lax.convert_element_type(x, some_extended_dtype) where x.dtype is not accepted, e.g. int -> float8_e4m3fn without going through the representation dtype, or a custom extended dtype whose rules disallow the pair.

Common situations: Assuming astype(float8) works from any numeric type; custom extended dtypes with restrictive allow_conversion; version changes in ml_dtypes rules.

Related errors

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

Appendix: source

Thrown at jax/_src/lax/lax.py:5503

  return trace.stage_value(args[0])

stage_p.def_bind_with_trace(_stage_bind_with_trace)


def _to_edtype_abstract_eval(x, *, edtype):
  assert (isinstance(edtype, dtypes.ExtendedDType) and
          not isinstance(x.dtype, dtypes.ExtendedDType))
  # For backward compatibility, if the edtype rules have a `convert_to` method,
  # use that rather than looking for an `allow_conversion: bool` attribute.
  if not isinstance(x, ShapedArray):
    raise TypeError("can only convert to an extended dtype on an array type,"
                    f"but got {type(x)}")
  if convert_to := getattr(edtype._rules, 'convert_to', None):
    allow_conversion = convert_to(x.dtype, edtype)
  else:
    allow_conversion = edtype._rules.allow_conversion
  if not allow_conversion:
    raise ValueError(
        f"Cannot convert_element_type from {dtype_to_string(x.dtype)} "
        f"to {dtype_to_string(edtype)}")
  rep_aval = core.physical_element_aval(edtype)
  assert tuple(rep_aval.sharding.spec) == (None,) * rep_aval.ndim
  if x.dtype != rep_aval.dtype:
    raise ValueError(
        "can only convert to extended dtype from its representation dtype, "
        f"but tried to convert from {dtype_to_string(x.dtype)} to "
        f"{dtype_to_string(edtype)} which doesn't match the representation type "
        f"{dtype_to_string(rep_aval.dtype)}.")
  if x.ndim < rep_aval.ndim:
    raise ValueError(
        "can only convert to extended dtype from an array of its "
        f"representation type, but the extended dtype {dtype_to_string(edtype)}"
        f" has a representation shape {rep_aval.shape} (rank {rep_aval.ndim}) "
        f"while the given representation array has shape {x.shape} (rank "
        f"{x.ndim} < {rep_aval.ndim}).")
  n = x.ndim - rep_aval.ndim

View on GitHub (pinned to 1e1c6a8fc0)