jax-ml/jax · error · ValueError

can only convert to extended dtype from an array with traili

Error message

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}

What it means

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.

Source

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

    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

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)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Shard only the leading (non-representation) axes; leave trailing axes unsharded (None in the PartitionSpec)
2. Convert to the extended dtype first, then reshard the result along logical axes
3. Reshape so representation components live on an unsharded axis
4. Use jax.lax.with_sharding_constraint after conversion rather than before

Example fix

// before
x = lax.with_sharding_constraint(flat, NamedSharding(mesh, P(None, 'i')))
out = lax.convert_element_type(x, edtype)

// after
x = lax.with_sharding_constraint(flat, NamedSharding(mesh, P('i', None)))
out = lax.convert_element_type(x, edtype)

n = x.ndim - physical_element_aval(edtype).ndim
assert all(s is None for s in x.sharding.spec[n:]), 'trailing axes must be unsharded'

When it happens

Trigger: 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,).

Common situations: 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.

Related errors

• Length of sharding.spec ({len(out_s.spec)}) must be equal to
• unreduced cannot contain None. All elements in unreduced sho
• reduced cannot contain None. All elements in reduced should
• `unreduced` and `reduced` argument to PartitionSpec cannot o
• Expected unreduced_kind to be of type `jax.sharding.Unreduce

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