{"record":{"id":"398b2ebf692f6a36","repo":"jax-ml/jax","slug":"can-only-convert-to-an-extended-dtype-on-an-array","errorCode":null,"errorMessage":"can only convert to an extended dtype on an array type,but got {type(x)}","messagePattern":"can only convert to an extended dtype on an array type,but got (.+?)","errorType":"exception","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/lax.py","lineNumber":5496,"sourceCode":"stage_p.def_impl(_stage_impl)\nbatching.defvectorized(stage_p)\nad.deflinear2(stage_p, lambda ct, _: [ct])\nmlir.register_lowering(stage_p, lambda ctx, operand: [operand])\npe.const_fold_rules[stage_p] = lambda consts, params, out_avals: consts\n\ndef _stage_bind_with_trace(trace, args, avals, params):\n  return trace.stage_value(args[0])\n\nstage_p.def_bind_with_trace(_stage_bind_with_trace)\n\n\ndef _to_edtype_abstract_eval(x, *, edtype):\n  assert (isinstance(edtype, dtypes.ExtendedDType) and\n          not isinstance(x.dtype, dtypes.ExtendedDType))\n  # For backward compatibility, if the edtype rules have a `convert_to` method,\n  # use that rather than looking for an `allow_conversion: bool` attribute.\n  if not isinstance(x, ShapedArray):\n    raise TypeError(\"can only convert to an extended dtype on an array type,\"\n                    f\"but got {type(x)}\")\n  if convert_to := getattr(edtype._rules, 'convert_to', None):\n    allow_conversion = convert_to(x.dtype, edtype)\n  else:\n    allow_conversion = edtype._rules.allow_conversion\n  if not allow_conversion:\n    raise ValueError(\n        f\"Cannot convert_element_type from {dtype_to_string(x.dtype)} \"\n        f\"to {dtype_to_string(edtype)}\")\n  rep_aval = core.physical_element_aval(edtype)\n  assert tuple(rep_aval.sharding.spec) == (None,) * rep_aval.ndim\n  if x.dtype != rep_aval.dtype:\n    raise ValueError(\n        \"can only convert to extended dtype from its representation dtype, \"\n        f\"but tried to convert from {dtype_to_string(x.dtype)} to \"\n        f\"{dtype_to_string(edtype)} which doesn't match the representation type \"\n        f\"{dtype_to_string(rep_aval.dtype)}.\")\n  if x.ndim < rep_aval.ndim:","sourceCodeStart":5478,"sourceCodeEnd":5514,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/lax.py#L5478-L5514","documentation":"_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.","triggerScenarios":"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.","commonSituations":"Using jax.experimental.export / dynamic shapes with extended dtypes; custom primitives returning non-shaped avals; partial-eval traces that produce non-ShapedArray avals.","solutions":["Avoid mixing symbolic/dynamic dimensions with extended dtype conversion; use static shapes","Check type(x) at trace time to detect dynamic-shape avals before converting","Update JAX: newer versions may support DShapedArray here","Convert to the representation dtype (e.g. uint8 for float8) instead, using bitcast where appropriate"],"exampleFix":null,"handlingStrategy":"type-guard","validationCode":"from jax.core import ShapedArray\nassert isinstance(aval, ShapedArray), 'requires ShapedArray aval'","typeGuard":"def is_shaped_aval(aval) -> bool:\n    from jax.core import ShapedArray\n    return isinstance(aval, ShapedArray)","tryCatchPattern":"try:\n    y = lax.convert_element_type(x, edtype)\nexcept TypeError:\n    raise RuntimeError('extended-dtype conversion needs static shapes') from None","preventionTips":["Avoid extended dtype conversion on dynamically shaped tracer paths","Pin static shapes in export pipelines","Test custom primitives' avals are ShapedArray"],"tags":["jax","extended-dtype","abstract-eval","dynamic-shape","typeerror"],"backgroundTag":"unsupported-aval-type","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}