jax-ml/jax · error · ValueError

Cannot convert_element_type from {dtype_to_string(x.dtype)}

Error message

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

What it means

Converting FROM an extended dtype to a regular dtype must be permitted by the dtype's rules (convert_from callback or allow_conversion flag). If the target dtype is not an allowed destination (e.g. converting float8 directly to int32 when rules disallow it), ValueError is raised.

Source

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

          convert_element_type(t, core.primal_dtype_to_tangent_dtype(edtype)))
ad.primitive_transposes[to_edtype_p] = \
    lambda ct, x, edtype: [from_edtype_p.bind(ct, dtype=x.aval.dtype)]
batching.defvectorized(to_edtype_p)
mlir.register_lowering(to_edtype_p, lambda _, x, **__: [x])


def _from_edtype_abstract_eval(x, *, dtype):
  assert (isinstance(x.dtype, dtypes.ExtendedDType) and
          not isinstance(dtype, dtypes.ExtendedDType))
  if not isinstance(x, ShapedArray):
    raise TypeError("can only convert from an extended dtype on an array type,"
                    f"but got {type(x)}")
  if convert_from := getattr(x.dtype._rules, 'convert_from', None):
    allow_conversion = convert_from(x.dtype, dtype)
  else:
    allow_conversion = x.dtype._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(dtype)}")
  rep_aval = core.physical_element_aval(x.dtype)
  assert tuple(rep_aval.sharding.spec) == (None,) * rep_aval.ndim
  if rep_aval.dtype != dtype:
    raise ValueError(
        "can only convert from extended dtype to its representation dtype, "
        f"but tried to convert from {dtype_to_string(x.dtype)} to "
        f"{dtype_to_string(dtype)} which doesn't match the representation type "
        f"{dtype_to_string(rep_aval.dtype)}.")
  if isinstance(x, ShapedArray):
    return x.update(shape=(*x.shape, *rep_aval.shape), dtype=dtype)
  else:
    assert False  # unreachable, see isinstance check above

from_edtype_p = Primitive('from_edtype')
from_edtype_p.def_impl(partial(dispatch.apply_primitive, from_edtype_p))
from_edtype_p.def_abstract_eval(_from_edtype_abstract_eval)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Convert to the representation dtype first (e.g. -> uint8), then cast to the desired type
For float8, go through float32 for numeric semantics
Inspect dtype._rules.convert_from to learn allowed targets
Extend the rules for custom extended dtypes where the conversion is meaningful

Example fix

// before
y = x_f8.astype(jnp.int32)  # disallowed pair

// after
y = x_f8.astype(jnp.float32).astype(jnp.int32)

Defensive patterns

Strategy: validation

Validate before calling

def from_conversion_allowed(edtype, target):
    fn = getattr(edtype._rules, 'convert_from', None)
    return fn(edtype, target) if fn else edtype._rules.allow_conversion

Try / catch

try:
    y = x_ed.astype(dt)
except ValueError:
    y = x_ed.astype(rep_dtype(x_ed.dtype)).astype(dt)

Prevention

Convert to representation dtype first, then cast
Whitelist target dtypes at API boundaries
Read the dtype rules before adding new cast targets

When it happens

Trigger: jax.lax.convert_element_type(x_edtype, dtype) where the pair (x.dtype, dtype) is rejected by the dtype rules, e.g. float8 -> int, or a custom extended dtype with restrictive convert_from.

Common situations: Trying to view/convert float8 payloads as integers directly; custom dtype ecosystems with explicit conversion whitelists; assuming numpy-style any-to-any casts.

Related errors

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