{"record":{"id":"35691353e8af4d2e","repo":"jax-ml/jax","slug":"can-only-convert-to-extended-dtype-from-an-array-w","errorCode":null,"errorMessage":"can only convert to extended dtype from an array with trailing axes that are not explicitly sharded, but tried to convert from {x.str_short(short_dtypes=True)} to an extended dtype with element shape {rep_aval.shape}","messagePattern":"can only convert to extended dtype from an array with trailing axes that are not explicitly sharded, but tried to convert from (.+?) to an extended dtype with element shape (.+?)","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/lax.py","lineNumber":5533,"sourceCode":"    raise ValueError(\n        \"can only convert to extended dtype from an array of its \"\n        f\"representation type, but the extended dtype {dtype_to_string(edtype)}\"\n        f\" has a representation shape {rep_aval.shape} (rank {rep_aval.ndim}) \"\n        f\"while the given representation array has shape {x.shape} (rank \"\n        f\"{x.ndim} < {rep_aval.ndim}).\")\n  n = x.ndim - rep_aval.ndim\n  shape_prefix, shape_suffix = x.shape[:n], x.shape[n:]\n  if shape_suffix != rep_aval.shape:\n    raise ValueError(\n        \"can only convert to extended dtype from an array of its \"\n        f\"representation type, but the extended dtype {dtype_to_string(edtype)}\"\n        f\" has a representation shape {rep_aval.shape} while the given \"\n        f\"representation array has shape {x.shape}, so the shape suffix \"\n        f\"does not match: given {shape_suffix} but required {rep_aval.shape}.\")\n  if isinstance(x, ShapedArray):\n    spec_prefix, spec_suffix = x.sharding.spec[:n], x.sharding.spec[n:]\n    if tuple(spec_suffix) != (None,) * len(spec_suffix):\n      raise ValueError(\n          \"can only convert to extended dtype from an array with trailing \"\n          \"axes that are not explicitly sharded, but tried to convert from \"\n          f\"{x.str_short(short_dtypes=True)} to an extended dtype with element \"\n          f\"shape {rep_aval.shape}\")\n    return x.update(shape=shape_prefix, dtype=edtype,\n                    sharding=x.sharding.update(spec=spec_prefix))\n  else:\n    assert False  # unreachable, see isinstance check above\n\nto_edtype_p = Primitive('to_edtype')\nto_edtype_p.def_impl(partial(dispatch.apply_primitive, to_edtype_p))\nto_edtype_p.def_abstract_eval(_to_edtype_abstract_eval)\nad.defjvp(to_edtype_p,\n          lambda t, x, edtype:\n          convert_element_type(t, core.primal_dtype_to_tangent_dtype(edtype)))\nad.primitive_transposes[to_edtype_p] = \\\n    lambda ct, x, edtype: [from_edtype_p.bind(ct, dtype=x.aval.dtype)]\nbatching.defvectorized(to_edtype_p)","sourceCodeStart":5515,"sourceCodeEnd":5551,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/lax.py#L5515-L5551","documentation":"When converting to an extended dtype with multi-element representation, the trailing axes of the input (those absorbed into each element) must not carry any explicit sharding in the input's sharding spec — an element cannot be split across devices. If any of the last len(spec_suffix) spec entries is not None, ValueError is raised.","triggerScenarios":"convert_element_type under a jit/pjit where the operand's NamedSharding shards one of the trailing representation axes, e.g. sharding=PartitionSpec(None, 'i') with representation shape (2,).","commonSituations":"Sharding the flattened representation buffer along its last axis before converting back to the extended dtype; migrating layouts between logical and physical forms inside sharded computations.","solutions":["Shard only the leading (non-representation) axes; leave trailing axes unsharded (None in the PartitionSpec)","Convert to the extended dtype first, then reshard the result along logical axes","Reshape so representation components live on an unsharded axis","Use jax.lax.with_sharding_constraint after conversion rather than before"],"exampleFix":"// before\nx = lax.with_sharding_constraint(flat, NamedSharding(mesh, P(None, 'i')))\nout = lax.convert_element_type(x, edtype)\n\n// after\nx = lax.with_sharding_constraint(flat, NamedSharding(mesh, P('i', None)))\nout = lax.convert_element_type(x, edtype)","handlingStrategy":"validation","validationCode":"n = x.ndim - physical_element_aval(edtype).ndim\nassert all(s is None for s in x.sharding.spec[n:]), 'trailing axes must be unsharded'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Never shard representation (trailing) axes of physical buffers","Convert to the extended dtype before constraining sharding on logical axes","Test pjit pipelines with multi-element-representation dtypes explicitly"],"tags":["jax","extended-dtype","sharding","partition-spec"],"backgroundTag":"sharding-spec-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}