jax-ml/jax · error · ValueError

can only convert to extended dtype from its representation d

Error message

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

What it means

Beyond dtype kind checks, converting TO an extended dtype requires the source dtype to exactly equal the extended dtype's representation dtype (e.g. uint8 for some 8-bit extension types). This error fires when the conversion was allowed by the rules but the actual source dtype differs from the physical representation type.

Source

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

  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
  shape_prefix, shape_suffix = x.shape[:n], x.shape[n:]
  if shape_suffix != rep_aval.shape:
    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} while the given "

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Match the source dtype exactly to the representation dtype (query it via jax.core.physical_element_aval(edtype).dtype)
Cast x to the representation dtype before the extended-dtype conversion
Verify signedness (uint8 vs int8) of the representation
For custom dtypes, document/expose the representation dtype in the rules

Example fix

// before
x_ed = lax.convert_element_type(jnp.uint16_array, edtype)  # rep is uint8

// after
x_ed = lax.convert_element_type(jnp.uint16_array.astype(jnp.uint8), edtype)

Defensive patterns

Strategy: validation

Validate before calling

from jax.core import physical_element_aval
rep = physical_element_aval(edtype).dtype
if x.dtype != rep:
    x = x.astype(rep)

Try / catch

try:
    out = lax.convert_element_type(x, edtype)
except ValueError:
    out = lax.convert_element_type(x.astype(rep_dtype(edtype)), edtype)

Prevention

Query and cache physical_element_aval(edtype).dtype for each extended dtype you use
Standardize buffers to the exact representation dtype (incl. signedness)
Validate at API boundaries where raw buffers enter

When it happens

Trigger: convert_element_type(x, edtype) where x.dtype != physical_element_aval(edtype).dtype, e.g. passing uint16 when the extended dtype is stored as uint8.

Common situations: Custom extended dtypes whose representation is not obvious; feeding buffers read with a different integer width; mixing jnp.uint8 and numpy uint8 vs int8 signedness mistakes.

Related errors

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