jax-ml/jax · error · ValueError

can only convert from extended dtype to its representation d

Error message

can only convert from extended dtype to its representation dtype, but tried to convert from {dtype_to_string(x.dtype)} to {dtype_to_string(dtype)} which doesn't match the representation type {dtype_to_string(rep_aval.dtype)}.

What it means

After rule-level approval, converting from an extended dtype requires the target dtype to exactly equal the extended dtype's representation dtype (e.g. uint8). If the requested dtype differs from the physical storage type, ValueError is raised with both dtypes and the required representation type.

Source

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

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)
ad.defjvp(from_edtype_p,
          lambda t, x, dtype:
          convert_element_type(t, core.primal_dtype_to_tangent_dtype(dtype)))
ad.primitive_transposes[from_edtype_p] = \
    lambda ct, x, dtype: [to_edtype_p.bind(ct, edtype=x.dtype)]
batching.defvectorized(from_edtype_p)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Use exactly the representation dtype: query it with jax.core.physical_element_aval(x.dtype).dtype
Then cast the result to the final desired dtype in a second step
Double-check uint vs int signedness of the representation
For numeric semantics, convert via the allowed path (edtype -> rep dtype -> float) per the rules

Example fix

// before
y = x_ed.astype(jnp.int8)  # representation is uint8

// after
y = x_ed.astype(jnp.uint8).astype(jnp.int8)

Defensive patterns

Strategy: validation

Validate before calling

from jax.core import physical_element_aval
rep = physical_element_aval(x.dtype).dtype
if target != rep:
    plan = (rep, target)  # two-step conversion

Try / catch

try:
    y = lax.convert_element_type(x_ed, target)
except ValueError:
    y = lax.convert_element_type(x_ed, rep_of(x_ed)).astype(target)

Prevention

Cache physical_element_aval(...).dtype per extended dtype
Check signedness (uint8 vs int8) of representation types
Write helper convert_from_ed(x, target) that does the two-step cast

When it happens

Trigger: convert_element_type(x_ed, dtype) with dtype != physical_element_aval(x_ed.dtype).dtype, e.g. reading a uint8-represented dtype out as int8 or float32.

Common situations: Signedness mixups (int8 vs uint8 representation); assuming the logical type (float8) can be materialized directly as float32; changing representation assumptions across versions of a custom dtype.

Related errors

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