jax-ml/jax · error · ValueError
`out_sharding` passed to `broadcast_in_dim` can only contain
Error message
`out_sharding` passed to `broadcast_in_dim` can only contain unreduced of kind `sum`. Got out_sharding={sharding} What it means
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.
Source
Thrown at jax/_src/lax/lax.py:6960
def _broadcast_in_dim_sharding_rule(operand, *, shape, broadcast_dimensions,
sharding):
if sharding is not None:
return sharding
bds = set(broadcast_dimensions)
orig_spec = iter(operand.sharding.spec.partitions)
new_spec = [next(orig_spec) if i in bds else None for i in range(len(shape))]
assert next(orig_spec, None) is None
mesh = (get_abstract_mesh() if operand.sharding.mesh.empty else
operand.sharding.mesh)
return operand.sharding.update(
mesh=mesh, spec=operand.sharding.spec.update(partitions=new_spec))
def _broadcast_in_dim_unreduced_rule(operand, sharding):
if sharding is not None and sharding.mesh.are_all_axes_explicit:
out = sharding.spec.unreduced
if out and sharding.spec.unreduced_kind is not UnreducedKind.sum:
raise ValueError(
'`out_sharding` passed to `broadcast_in_dim` can only contain'
f' unreduced of kind `sum`. Got out_sharding={sharding}')
else:
out = getu(operand)
kind = UnreducedKind.sum if out else None
return out, kind
def _broadcast_in_dim_reduced_rule(operand, sharding):
if sharding is not None and sharding.mesh.are_all_axes_explicit:
return sharding.spec.reduced
return getr(operand)
def _broadcast_in_dim_ur_rule(operand, *, shape, broadcast_dimensions, sharding):
out_unreduced, kind = _broadcast_in_dim_unreduced_rule(operand, sharding)
out_reduced = _broadcast_in_dim_reduced_rule(operand, sharding)
return out_unreduced, out_reduced, kind
def _broadcast_in_dim_memory_space_rule(operand, *, shape, broadcast_dimensions,View on GitHub (pinned to 1e1c6a8fc0)
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
Example fix
// before out = lax.broadcast_in_dim(x, shape, bd, sharding=bad_sharding) # non-sum unreduced // after out = lax.broadcast_in_dim(x, shape, bd, sharding=None)
Defensive patterns
Strategy: validation
Validate before calling
sh = out_sharding assert sh is None or not sh.spec.unreduced or sh.spec.unreduced_kind is UnreducedKind.sum
Prevention
- Don't reuse reduction-output shardings on broadcasts
- Pass sharding=None unless you constructed it deliberately
When it happens
Trigger: 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.
Common situations: 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.
Related errors
- unreduced rule for {name} is not implemented. Please file an
- {name} cannot accept args which are unreduced. Got {a.str_sh
- Query, key and value should have same sharding.
- 0th dimension of leaf passed to `jax.lax.map` should be repl
- The denominator cannot be unreduced passed to `div`. Got {y=
AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27).
Data as JSON: /api/errors/b3f013b321418189.
Report an issue: GitHub.