jax-ml/jax · error · ValueError

out_specs passed to shard_map should be equal to the reduced

Error message

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

What it means

When checking is enabled, out_specs' reduced tuple must equal the output aval's reduced tuple; the declared output reduction must agree with the reduction the body actually produced.

Source

Thrown at jax/_src/shard_map.py:899

  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):
      if name_s and not aval_s:
        out_spec.append(name_s)
      elif aval_s and not name_s:
        out_spec.append(aval_s)
      elif not name_s and not aval_s:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Declare the same reduced axes in out_specs as on the output aval
2. Fully reduce the output inside the body if plain out_specs is desired
3. Log aval.str_short(True) of outputs during development to align specs

Example fix

// before
out = shard_map(f, mesh, x, out_specs=P('d'))
// after
out = shard_map(f, mesh, x, out_specs=P('d', reduced=('y',)))

try: shard_map(...) except ValueError as e: if 'reduced present' in str(e): add the reduced axes to out_specs; else: raise

When it happens

Trigger: Body yields arrays with reduced=('y',) but out_specs declares reduced=() or different axes.

Common situations: Partial-manual pipelines returning partially reduced tensors (gradient accumulation patterns) with plain out_specs.

Related errors

• in_specs containing reduced {spec} passed to shard_map shoul
• out_specs passed to shard_map should be equal to the unreduc
• out_specs passed to shard_map should be equal to the unreduc
• 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/f8d8427af2303693. Report an issue: GitHub.