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} while the given representation array has shape {x.shape}, so the shape suffix does not match: given {shape_suffix} but required {rep_aval.shape}.

What it means

For extended dtypes with multi-element representation, the trailing shape of the input must exactly equal the representation shape. After peeling off the leading dims, if the suffix (x.shape[n:]) != rep_aval.shape, the array cannot be reinterpreted as elements of that extended dtype.

Source

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

  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(
          "can only convert to extended dtype from an array with trailing "
          "axes that are not explicitly sharded, but tried to convert from "
          f"{x.str_short(short_dtypes=True)} to an extended dtype with element "
          f"shape {rep_aval.shape}")
    return x.update(shape=shape_prefix, dtype=edtype,
                    sharding=x.sharding.update(spec=spec_prefix))
  else:
    assert False  # unreachable, see isinstance check above

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Reshape x so its trailing dims equal rep_aval.shape exactly
2. Verify the required suffix via jax.core.physical_element_aval(edtype).shape
3. Reorder/transpose if the representation expects components along a different axis
4. Fix upstream shape bugs (wrong axis length) before conversion

Example fix

// before
out = lax.convert_element_type(x.reshape(2, 5), edtype)  # needs suffix (2,)

// after
out = lax.convert_element_type(x.reshape(5, 2), edtype)

from jax.core import physical_element_aval
rep = physical_element_aval(edtype)
n = x.ndim - rep.ndim
assert n >= 0 and x.shape[n:] == rep.shape, (x.shape, rep.shape)

try:
    out = lax.convert_element_type(x, edtype)
except ValueError:
    out = lax.convert_element_type(x.reshape(*x.shape[:-rep.ndim], *rep.shape), edtype)

When it happens

Trigger: convert_element_type(x, edtype) where x.shape[-k:] doesn't match the representation shape, e.g. shape (2, 5) with representation (2,) — the last dim must be 2 but is 5.

Common situations: Mismatched flattened buffer sizes for composite dtypes; off-by-one reshapes before conversion; representation ordering confusion (e.g. (2,) vs (1,2)).

Related errors

• can only convert to extended dtype from an array of its repr
• scan got `length` argument of {} which disagrees with leadin
• conv_general_dilated batch_group_count must divide lhs batch
• conv_general_dilated rhs output feature dimension size must
• conv_general_dilated window and window_strides must have the

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