jax-ml/jax · error · TypeError

can only convert to an extended dtype on an array type,but g

Error message

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

What it means

_to_edtype_abstract_eval handles converting a normal array aval into an extended dtype (e.g. float8, custom extension types). The conversion logic requires a ShapedArray with a concrete shape; receiving another aval type (like a DShapedArray with symbolic dimensions, or a token/special aval) is unsupported, so TypeError is raised.

Source

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

stage_p.def_impl(_stage_impl)
batching.defvectorized(stage_p)
ad.deflinear2(stage_p, lambda ct, _: [ct])
mlir.register_lowering(stage_p, lambda ctx, operand: [operand])
pe.const_fold_rules[stage_p] = lambda consts, params, out_avals: consts

def _stage_bind_with_trace(trace, args, avals, params):
  return trace.stage_value(args[0])

stage_p.def_bind_with_trace(_stage_bind_with_trace)


def _to_edtype_abstract_eval(x, *, edtype):
  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:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Avoid mixing symbolic/dynamic dimensions with extended dtype conversion; use static shapes
2. Check type(x) at trace time to detect dynamic-shape avals before converting
3. Update JAX: newer versions may support DShapedArray here
4. Convert to the representation dtype (e.g. uint8 for float8) instead, using bitcast where appropriate

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

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

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

When it happens

Trigger: Calling convert_element_type with new_dtype being an ExtendedDtype on an operand whose aval is not a ShapedArray, e.g. arrays with dynamic/symbolic shape dimensions under export or shape polymorphism.

Common situations: Using jax.experimental.export / dynamic shapes with extended dtypes; custom primitives returning non-shaped avals; partial-eval traces that produce non-ShapedArray avals.

Related errors

• can only convert from an extended dtype on an array type,but
• 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/398b2ebf692f6a36. Report an issue: GitHub.