{"record":{"id":"b3f013b321418189","repo":"jax-ml/jax","slug":"out-sharding-passed-to-broadcast-in-dim-can-on","errorCode":null,"errorMessage":"`out_sharding` passed to `broadcast_in_dim` can only contain unreduced of kind `sum`. Got out_sharding={sharding}","messagePattern":"`out_sharding` passed to `broadcast_in_dim` can only contain unreduced of kind `sum`\\. Got out_sharding=(.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/lax.py","lineNumber":6960,"sourceCode":"\ndef _broadcast_in_dim_sharding_rule(operand, *, shape, broadcast_dimensions,\n                                    sharding):\n  if sharding is not None:\n    return sharding\n  bds = set(broadcast_dimensions)\n  orig_spec = iter(operand.sharding.spec.partitions)\n  new_spec = [next(orig_spec) if i in bds else None for i in range(len(shape))]\n  assert next(orig_spec, None) is None\n  mesh = (get_abstract_mesh() if operand.sharding.mesh.empty else\n          operand.sharding.mesh)\n  return operand.sharding.update(\n      mesh=mesh, spec=operand.sharding.spec.update(partitions=new_spec))\n\ndef _broadcast_in_dim_unreduced_rule(operand, sharding):\n  if sharding is not None and sharding.mesh.are_all_axes_explicit:\n    out = sharding.spec.unreduced\n    if out and sharding.spec.unreduced_kind is not UnreducedKind.sum:\n      raise ValueError(\n          '`out_sharding` passed to `broadcast_in_dim` can only contain'\n          f' unreduced of kind `sum`. Got out_sharding={sharding}')\n  else:\n    out = getu(operand)\n  kind = UnreducedKind.sum if out else None\n  return out, kind\n\ndef _broadcast_in_dim_reduced_rule(operand, sharding):\n  if sharding is not None and sharding.mesh.are_all_axes_explicit:\n    return sharding.spec.reduced\n  return getr(operand)\n\ndef _broadcast_in_dim_ur_rule(operand, *, shape, broadcast_dimensions, sharding):\n  out_unreduced, kind = _broadcast_in_dim_unreduced_rule(operand, sharding)\n  out_reduced = _broadcast_in_dim_reduced_rule(operand, sharding)\n  return out_unreduced, out_reduced, kind\n\ndef _broadcast_in_dim_memory_space_rule(operand, *, shape, broadcast_dimensions,","sourceCodeStart":6942,"sourceCodeEnd":6978,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/lax.py#L6942-L6978","documentation":"In GSPMD-style sharding propagation, an out_sharding handed to broadcast_in_dim may only carry unreduced axes of kind 'sum'. Other unreduced kinds (e.g. from a different reduction semantics) are rejected because broadcasting can't preserve them.","triggerScenarios":"Calling broadcast_in_dim with an out_sharding whose spec contains unreduced axes with a kind other than UnreducedKind.sum, on a mesh with all-explicit axes.","commonSituations":"Manual GSPMD/sharding-annotation work where a NamedSharding with unreduced specs built for a reduction output is reused on a broadcast; internal/migrating code after JAX sharding-spec API changes.","solutions":["Strip or fix the unreduced axes in the out_sharding (only 'sum' kind is allowed) before passing it","Pass sharding=None and let it be inferred from the operand","Regenerate the sharding from the intended output rather than reusing one from a reduction"],"exampleFix":"// before\nout = lax.broadcast_in_dim(x, shape, bd, sharding=bad_sharding)  # non-sum unreduced\n// after\nout = lax.broadcast_in_dim(x, shape, bd, sharding=None)","handlingStrategy":"validation","validationCode":"sh = out_sharding\nassert sh is None or not sh.spec.unreduced or sh.spec.unreduced_kind is UnreducedKind.sum","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Don't reuse reduction-output shardings on broadcasts","Pass sharding=None unless you constructed it deliberately"],"tags":["jax","broadcast-in-dim","sharding","gspmd","unreduced"],"backgroundTag":"invalid-sharding-specification","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}