jax-ml/jax · error · ValueError

out_specs passed to shard_map should be equal to the unreduc

Error message

out_specs passed to shard_map should be equal to the unreduced_kind present on the out_aval. Got out_specs={spec} and out_aval={aval.str_short(True)}

What it means

When checking is on, shard_map requires out_specs' unreduced_kind to equal the output aval's unreduced_kind. The combination rule declared for replicated partial outputs must match what the computation actually produced.

Source

Thrown at jax/_src/shard_map.py:894

  unreduced = aval.sharding.spec.unreduced if check_vma else frozenset()
  reduced = aval.sharding.spec.reduced if check_vma else frozenset()
  u_kind = aval.sharding.spec.unreduced_kind if check_vma else None
  mat = core.ManualAxisType(varying=vma, unreduced=unreduced, reduced=reduced,
                            unreduced_kind=u_kind)
  return aval.update(shape=new_shape, sharding=new_sharding,
                     manual_axis_type=mat)
core.shard_aval_handlers[core.ShapedArray] = _shard_shaped_array

def _unshard_shaped_array(mesh: Mesh, check_vma, spec, aval: core.ShapedArray
                          ) -> core.ShapedArray:
  assert isinstance(aval, core.ShapedArray)
  if check_vma and spec.unreduced != aval.mat.unreduced:
    raise ValueError(
        "out_specs passed to shard_map should be equal to the unreduced"
        f" present on the out_aval. Got out_specs={spec} and"
        f" out_aval={aval.str_short(True)}")
  if check_vma and spec.unreduced_kind is not aval.mat.unreduced_kind:
    raise ValueError(
        "out_specs passed to shard_map should be equal to the unreduced_kind"
        f" present on the out_aval. Got out_specs={spec} and"
        f" out_aval={aval.str_short(True)}")
  if check_vma and spec.reduced != aval.mat.reduced:
    raise ValueError(
        "out_specs passed to shard_map should be equal to the reduced present"
        f" on the out_aval. Got out_specs={spec} and"
        f" out_aval={aval.str_short(True)}")
  names = _spec_to_names(spec)
  new_shape = tuple(sz * prod(mesh.shape[n] for n in names.get(i, ()))
                    for i, sz in enumerate(aval.shape))
  names_spec = spec._normalized_spec_for_aval(aval.ndim).partitions
  if aval.ndim == 0:
    out_spec = P(unreduced=spec.unreduced, reduced=spec.reduced,
                 unreduced_kind=spec.unreduced_kind)
  else:
    out_spec = []
    for name_s, aval_s in zip(names_spec, aval.sharding.spec.partitions):

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Match unreduced_kind in out_specs to the output aval's value
2. Regenerate outputs with the intended kind by fixing upstream shardings/ops
3. Inspect aval.mat.unreduced_kind to confirm the expected value

Example fix

// before
out = shard_map(f, mesh, x, out_specs=P(unreduced=('r',), unreduced_kind='mul'))
// after
out = shard_map(f, mesh, x, out_specs=P(unreduced=('r',), unreduced_kind='add'))

try: shard_map(...) except ValueError as e: if 'unreduced_kind' in str(e): align kind in out_specs; else: raise

When it happens

Trigger: Body returns arrays with unreduced_kind 'add' but out_specs declares 'mul' (or defaults differ).

Common situations: Switching reduction kind when porting code between JAX versions whose defaults for unreduced_kind changed.

Related errors

• in_specs containing unreduced_kind {spec} passed to shard_ma
• out_specs passed to shard_map should be equal to the unreduc
• out_specs passed to shard_map should be equal to the reduced
• shard_map out_specs rank error (msg from _spec_rank_error)
• shard_map out_specs vma error (msg from _inout_vma_error)

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