jax-ml/jax · error · ValueError

can only convert to extended dtype from an array of its repr

Error message

can only convert to extended dtype from an array of its representation type, but the extended dtype {dtype_to_string(edtype)} has a representation shape {rep_aval.shape} (rank {rep_aval.ndim}) while the given representation array has shape {x.shape} (rank {x.ndim} < {rep_aval.ndim}).

What it means

Some extended dtypes have multi-element representation shapes (an element occupies several physical elements, e.g. a complex-of-float8 stored as 2 uint8s). To convert into such a dtype, the input array must have at least as many trailing dimensions as the representation rank; if x.ndim < rep_aval.ndim, ValueError is raised.

Source

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

                    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 "
        f"representation array has shape {x.shape}, so the shape suffix "
        f"does not match: given {shape_suffix} but required {rep_aval.shape}.")
  if isinstance(x, ShapedArray):
    spec_prefix, spec_suffix = x.sharding.spec[:n], x.sharding.spec[n:]
    if tuple(spec_suffix) != (None,) * len(spec_suffix):
      raise ValueError(

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

  1. Ensure the input has at least rep_aval.ndim trailing dimensions (append axes with [...] if semantics allow)
  2. Reshape so the last dims match the representation shape
  3. Check core.physical_element_aval(edtype).ndim to know the required rank
  4. Split into per-component conversions if the composite conversion is not needed

Example fix

// before
out = lax.convert_element_type(scalar_x, composite_edtype)  # rep rank 2

// after
out = lax.convert_element_type(scalar_x[None, None], composite_edtype)
Defensive patterns

Strategy: validation

Validate before calling

from jax.core import physical_element_aval
rep = physical_element_aval(edtype)
assert x.ndim >= rep.ndim, f'need rank >= {rep.ndim}'

Try / catch

try:
    out = lax.convert_element_type(x, edtype)
except ValueError:
    out = lax.convert_element_type(x[..., None, None], edtype)  # pad to rep rank

Prevention

When it happens

Trigger: convert_element_type on a scalar or low-rank array to an extended dtype whose representation shape has rank > 0, e.g. shape () or (4,) into an edtype with representation shape (2,).

Common situations: Working with custom composite extended dtypes (e.g. pair/triple representations); assuming extended elements are always scalar-represented; flattening arrays before conversion.

Related errors


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