jax-ml/jax · error · ValueError

can only convert to extended dtype from its representation…

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.

Solutions

  1. Match the source dtype exactly to the representation dtype (query it via jax.core.physical_element_aval(edtype).dtype)
  2. Cast x to the representation dtype before the extended-dtype conversion
  3. Verify signedness (uint8 vs int8) of the representation
  4. 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

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.

Appendix: 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)