{"record":{"id":"63104a1b3ba4c16e","repo":"jax-ml/jax","slug":"out-sharding-passed-to-name-cannot-be-reduced-g","errorCode":null,"errorMessage":"out_sharding passed to {name} cannot be reduced. Got {out_sharding=}","messagePattern":"out_sharding passed to (.+?) cannot be reduced\\. Got (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/lax.py","lineNumber":8571,"sourceCode":"    axes = frozenset(axes)\n    used_spec = frozenset(\n        s for i, spec in enumerate(operand.sharding.spec.partitions)\n        if i in axes for s in (spec if isinstance(spec, tuple) else (spec,))\n    ) | operand.sharding.spec.unreduced\n    if not all(u in used_spec for u in out_sharding.spec.unreduced):\n      raise core.ShardingTypeError(\n          \"out_sharding's unreduced axes should be in operand's specs that\"\n          f' were {name} over. Got {operand=}, {axes=},'\n          f' unreduced_spec={out_sharding.spec.unreduced}')\n    out_u = out_sharding.spec.unreduced\n  else:\n    # TODO(yashkatariya): For max/min, do getu(operand, out_kind) and add tests\n    out_u = getu(operand)\n  return out_u, out_kind if out_u else None\n\ndef _reduce_op_reduced_rule(operand, out_sharding, name):\n  if out_sharding is not None and out_sharding.spec.reduced:\n    raise ValueError(\n        f'out_sharding passed to {name} cannot be reduced. Got {out_sharding=}')\n  return getr(operand)\n\ndef _reduce_sum_ur_rule(operand, *, axes, out_sharding):\n  out_unreduced, kind = _reduce_op_unreduced_rule(\n      operand, axes, out_sharding, UnreducedKind.sum, 'reduce_sum')\n  out_reduced = _reduce_op_reduced_rule(operand, out_sharding, 'reduce_sum')\n  return out_unreduced, out_reduced, kind\n\ndef _reduce_sum_dtype_rule(operand, *, axes, **_):\n  dt = _reduce_number_dtype_rule('reduce_sum', operand)\n  if (operand.dtype in [np.float16, dtypes.bfloat16] and\n      not config.allow_f16_reductions.value and\n      not all(core.definitely_equal(operand.shape[d], 1) for d in axes)):\n    raise ValueError(f\"reduce_sum on operand {operand.str_short(True)} is not \"\n                     \"allowed when jax_allow_f16_reductions=False.\")\n  return dt\n","sourceCodeStart":8553,"sourceCodeEnd":8589,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/lax.py#L8553-L8589","documentation":"The out_sharding (NamedSharding) passed to a reduction op like reduce_sum has a 'reduced' spec on some axis, meaning the caller asked for the output to be sharded along an axis that is being reduced away. That layout is impossible, so JAX rejects it during abstract evaluation.","triggerScenarios":"Calling reduce_sum(operand, axes, out_sharding=NamedSharding(mesh, P('data', None))) where the 'data' mesh axis maps to a dimension being reduced. Happens with the out_sharding keyword introduced for explicit output-sharding control of reductions.","commonSituations":"Migrating multi-host/multi-GPU code to the out_sharding API and reusing an input sharding for the output; forgetting that reduced dimensions must be replicated (None) in the output sharding spec.","solutions":["Set the reduced dimension's entry to None (replicated) in the out_sharding spec: P(None, 'data') instead of P('data', 'data') when reducing axis 0.","Compute the output sharding from the output shape, not the input shape.","If automatic sharding propagation is fine, omit out_sharding entirely."],"exampleFix":"# before\nout_sharding = NamedSharding(mesh, P('data', None))  # reducing axis 0\nlax.reduce_sum(x, (0,), out_sharding=out_sharding)\n# after\nout_sharding = NamedSharding(mesh, P(None, 'data'))\nlax.reduce_sum(x, (0,), out_sharding=out_sharding)","handlingStrategy":"validation","validationCode":"from jax.sharding import NamedSharding, PartitionSpec as P\nout_pspec = tuple(None if i in axes else ps[i] for i, ps in enumerate(in_pspec))\nout_sharding = NamedSharding(mesh, P(*out_pspec))","typeGuard":"def sharding_ok_for_reduction(axes, pspec):\n    return all(pspec[i] is None for i in axes)","tryCatchPattern":null,"preventionTips":["Derive output shardings from output shape, never reuse input specs blindly.","Remember reduced dims must be replicated (None) in out_sharding."],"tags":["jax","sharding","named-sharding","reduction","distributed"],"backgroundTag":"invalid-sharding-spec","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}