jax-ml/jax · error · TypeError

can only convert from an extended dtype on an array type,but

Error message

can only convert from an extended dtype on an array type,but got {type(x)}

What it means

Mirror of error 903 for the reverse direction: converting FROM an extended dtype requires the input aval to be a ShapedArray. Other aval types (dynamic-shape arrays, tokens, etc.) cannot carry the shape/dtype update needed for the conversion, so TypeError is raised.

Source

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

    assert False  # unreachable, see isinstance check above

to_edtype_p = Primitive('to_edtype')
to_edtype_p.def_impl(partial(dispatch.apply_primitive, to_edtype_p))
to_edtype_p.def_abstract_eval(_to_edtype_abstract_eval)
ad.defjvp(to_edtype_p,
          lambda t, x, edtype:
          convert_element_type(t, core.primal_dtype_to_tangent_dtype(edtype)))
ad.primitive_transposes[to_edtype_p] = \
    lambda ct, x, edtype: [from_edtype_p.bind(ct, dtype=x.aval.dtype)]
batching.defvectorized(to_edtype_p)
mlir.register_lowering(to_edtype_p, lambda _, x, **__: [x])


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

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Use static shapes for the operand when converting out of an extended dtype
2. Check isinstance(aval, ShapedArray) before issuing the conversion in meta-code
3. Update JAX in case newer releases support DShapedArray here
4. Convert via the physical representation using bitcast/reshape instead

from jax.core import ShapedArray
assert isinstance(aval, ShapedArray)

def is_shaped_aval(aval) -> bool:
    from jax.core import ShapedArray
    return isinstance(aval, ShapedArray)

try:
    y = lax.convert_element_type(x_ed, dt)
except TypeError:
    raise RuntimeError('conversion from extended dtype needs static shapes') from None

When it happens

Trigger: convert_element_type(x_with_edtype, standard_dtype) where x's aval is not ShapedArray, typically under shape polymorphism/export with symbolic dims.

Common situations: jax.experimental.export pipelines with extended dtypes; custom primitives with non-standard avals; tracing internals producing abstract avals without concrete shapes.

Related errors

• can only convert to an extended dtype on an array type,but g
• iteration over a 0-d array
• Formatting arguments to checkify.check need to be PyTrees of
• check_error takes an Error as argument, got type {type(error
• bool() not supported for instances of type '{0}' (did you me

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