{"record":{"id":"8595a676e385a30e","repo":"jax-ml/jax","slug":"sharding-passed-to-convert-element-type-can-on","errorCode":null,"errorMessage":"`sharding` passed to `convert_element_type` can only contain unreduced of kind `sum`. Got sharding={sharding}","messagePattern":"`sharding` passed to `convert_element_type` can only contain unreduced of kind `sum`\\. Got sharding=(.+?)","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/lax.py","lineNumber":5313,"sourceCode":"  return operand.shape\n\ndef _convert_element_type_sharding_rule(operand, *, new_dtype, weak_type,\n                                        sharding):\n  if sharding is None:\n    return operand.sharding\n  if sharding._is_concrete:\n    if isinstance(sharding, NamedSharding):\n      return NamedSharding(sharding.mesh.abstract_mesh, sharding.spec)\n    else:\n      return core.get_cur_mesh_sharding()\n  return sharding\n\ndef _convert_element_type_ur_rule(operand, *, new_dtype, weak_type, sharding):\n  if (sharding is not None and isinstance(sharding, NamedSharding) and\n      sharding.spec.unreduced):\n    kind = sharding.spec.unreduced_kind\n    if kind is not None and kind is not UnreducedKind.sum:\n      raise ValueError(\n          '`sharding` passed to `convert_element_type` can only contain'\n          f' unreduced of kind `sum`. Got sharding={sharding}')\n    unreduced = sharding.spec.unreduced\n  else:\n    unreduced = getu(operand)\n  reduced = (sharding.spec.reduced\n             if sharding is not None and isinstance(sharding, NamedSharding)\n             and sharding.spec.reduced else getr(operand))\n  kind = UnreducedKind.sum if unreduced else None\n  return unreduced, reduced, kind\n\ndef _convert_element_type_dtype_rule(operand, *, new_dtype, weak_type,\n                                     sharding):\n  return new_dtype\n\ndef _convert_element_type_weak_type_rule(operand, *, new_dtype, weak_type,\n                                         sharding):\n  return weak_type","sourceCodeStart":5295,"sourceCodeEnd":5331,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/lax.py#L5295-L5331","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","solutions":["Reduce (all_gather) the min/max-unreduced operand before converting its dtype","Pass sharding=None so the rule falls back to the operand's own unreduced set","Restructure so the dtype conversion happens before the min/max-style partial reduction","Build the sharding spec with UnreducedKind.sum if sums are actually intended"],"exampleFix":"// before\ny = lax.convert_element_type(x_min_unreduced, jnp.float32)\n\n// after\nx_full = lax.all_gather(x_min_unreduced, 'i')  # finish the reduction first\ny = lax.convert_element_type(x_full, jnp.float32)","handlingStrategy":"validation","validationCode":"def can_cast_unreduced(x, sharding):\n    if sharding is not None and getattr(sharding.spec, 'unreduced', False):\n        return sharding.spec.unreduced_kind in (None, UnreducedKind.sum)\n    return True","typeGuard":null,"tryCatchPattern":"try:\n    y = lax.convert_element_type(x, dt, sharding=s)\nexcept ValueError:\n    x = lax.all_gather(x, 'i')\n    y = lax.convert_element_type(x, dt)","preventionTips":["Finish min/max-style partial reductions before dtype conversion","Don't pass explicit sharding with non-sum unreduced kinds to convert_element_type","Unit-test sharded pipelines with each reduction kind you use"],"tags":["jax","sharding","dtype-conversion","lax"],"backgroundTag":"sharding-spec-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}