jax-ml/jax · error · ValueError

`sharding` passed to `convert_element_type` can only contain

Error message

`sharding` passed to `convert_element_type` can only contain unreduced of kind `sum`. Got sharding={sharding}

What it means

convert_element_type participates in the 'unreduced' sharding propagation system, but a dtype conversion is only valid across unreduced axes of kind 'sum' (the values are partial sums). If the sharding spec marks axes as unreduced with a different kind (e.g. 'min'/'max' from a different reduction), converting the element type would be semantically wrong, so ValueError is raised.

Source

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

  return operand.shape

def _convert_element_type_sharding_rule(operand, *, new_dtype, weak_type,
                                        sharding):
  if sharding is None:
    return operand.sharding
  if sharding._is_concrete:
    if isinstance(sharding, NamedSharding):
      return NamedSharding(sharding.mesh.abstract_mesh, sharding.spec)
    else:
      return core.get_cur_mesh_sharding()
  return sharding

def _convert_element_type_ur_rule(operand, *, new_dtype, weak_type, sharding):
  if (sharding is not None and isinstance(sharding, NamedSharding) and
      sharding.spec.unreduced):
    kind = sharding.spec.unreduced_kind
    if kind is not None and kind is not UnreducedKind.sum:
      raise ValueError(
          '`sharding` passed to `convert_element_type` can only contain'
          f' unreduced of kind `sum`. Got sharding={sharding}')
    unreduced = sharding.spec.unreduced
  else:
    unreduced = getu(operand)
  reduced = (sharding.spec.reduced
             if sharding is not None and isinstance(sharding, NamedSharding)
             and sharding.spec.reduced else getr(operand))
  kind = UnreducedKind.sum if unreduced else None
  return unreduced, reduced, kind

def _convert_element_type_dtype_rule(operand, *, new_dtype, weak_type,
                                     sharding):
  return new_dtype

def _convert_element_type_weak_type_rule(operand, *, new_dtype, weak_type,
                                         sharding):
  return weak_type

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Reduce (all_gather) the min/max-unreduced operand before converting its dtype
2. Pass sharding=None so the rule falls back to the operand's own unreduced set
3. Restructure so the dtype conversion happens before the min/max-style partial reduction
4. Build the sharding spec with UnreducedKind.sum if sums are actually intended

Example fix

// before
y = lax.convert_element_type(x_min_unreduced, jnp.float32)

// after
x_full = lax.all_gather(x_min_unreduced, 'i')  # finish the reduction first
y = lax.convert_element_type(x_full, jnp.float32)

def can_cast_unreduced(x, sharding):
    if sharding is not None and getattr(sharding.spec, 'unreduced', False):
        return sharding.spec.unreduced_kind in (None, UnreducedKind.sum)
    return True

try:
    y = lax.convert_element_type(x, dt, sharding=s)
except ValueError:
    x = lax.all_gather(x, 'i')
    y = lax.convert_element_type(x, dt)

When it happens

Trigger: Calling jax.lax.convert_element_type (or .astype under the hood in jit with sharding propagation) on an operand whose NamedSharding spec has unreduced_kind other than UnreducedKind.sum, or explicitly passing sharding= with such a spec.

Common situations: Composing automatic partially-reduced collectives (pmin/pmax style) followed by astype; manually constructing NamedSharding specs with unreduced kinds; version upgrades that introduced unreduced_kind semantics.

Related errors

• Denominator should be reduced along the same axes numerator
• Mapped away dimension of inputs passed to vmap should be sha
• Unmapped values passed to vmap cannot be sharded along the m
• to_dlpack can only pack a dlpack tensor from an array on a s
• __dlpack__ only supported for unsharded arrays.

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