jax-ml/jax · error · ValueError

out_sharding passed to {name} cannot be reduced. Got {out_sh

Error message

out_sharding passed to {name} cannot be reduced. Got {out_sharding=}

What it means

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.

Source

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

    axes = frozenset(axes)
    used_spec = frozenset(
        s for i, spec in enumerate(operand.sharding.spec.partitions)
        if i in axes for s in (spec if isinstance(spec, tuple) else (spec,))
    ) | operand.sharding.spec.unreduced
    if not all(u in used_spec for u in out_sharding.spec.unreduced):
      raise core.ShardingTypeError(
          "out_sharding's unreduced axes should be in operand's specs that"
          f' were {name} over. Got {operand=}, {axes=},'
          f' unreduced_spec={out_sharding.spec.unreduced}')
    out_u = out_sharding.spec.unreduced
  else:
    # TODO(yashkatariya): For max/min, do getu(operand, out_kind) and add tests
    out_u = getu(operand)
  return out_u, out_kind if out_u else None

def _reduce_op_reduced_rule(operand, out_sharding, name):
  if out_sharding is not None and out_sharding.spec.reduced:
    raise ValueError(
        f'out_sharding passed to {name} cannot be reduced. Got {out_sharding=}')
  return getr(operand)

def _reduce_sum_ur_rule(operand, *, axes, out_sharding):
  out_unreduced, kind = _reduce_op_unreduced_rule(
      operand, axes, out_sharding, UnreducedKind.sum, 'reduce_sum')
  out_reduced = _reduce_op_reduced_rule(operand, out_sharding, 'reduce_sum')
  return out_unreduced, out_reduced, kind

def _reduce_sum_dtype_rule(operand, *, axes, **_):
  dt = _reduce_number_dtype_rule('reduce_sum', operand)
  if (operand.dtype in [np.float16, dtypes.bfloat16] and
      not config.allow_f16_reductions.value and
      not all(core.definitely_equal(operand.shape[d], 1) for d in axes)):
    raise ValueError(f"reduce_sum on operand {operand.str_short(True)} is not "
                     "allowed when jax_allow_f16_reductions=False.")
  return dt

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. 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.
2. Compute the output sharding from the output shape, not the input shape.
3. If automatic sharding propagation is fine, omit out_sharding entirely.

Example fix

# before
out_sharding = NamedSharding(mesh, P('data', None))  # reducing axis 0
lax.reduce_sum(x, (0,), out_sharding=out_sharding)
# after
out_sharding = NamedSharding(mesh, P(None, 'data'))
lax.reduce_sum(x, (0,), out_sharding=out_sharding)

from jax.sharding import NamedSharding, PartitionSpec as P
out_pspec = tuple(None if i in axes else ps[i] for i, ps in enumerate(in_pspec))
out_sharding = NamedSharding(mesh, P(*out_pspec))

def sharding_ok_for_reduction(axes, pspec):
    return all(pspec[i] is None for i in axes)

When it happens

Trigger: 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.

Common situations: 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.

Related errors

• Reduced axes can only refer to mesh axes that is of type `Ex
• Mapped away dimension of inputs passed to vmap should be sha
• The 'sharding' attribute is not available on {self._error_re
• The is_fully_addressable property was called on {self._error
• Sharding spec {spec} implies that array axis {dim} is partit

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