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 present on the out_aval. Got out_specs={spec} and out_aval={aval.str_short(True)}

What it means

On the output side (when VMA checking is enabled), shard_map verifies that out_specs' unreduced tuple matches the unreduced information on the output aval produced by the mapped function. A mismatch means the declared output replication contradicts what the body computed.

Source

Thrown at jax/_src/shard_map.py:889

  manual_mesh = _as_manual_mesh(mesh, manual_axes)
  new_sharding = aval.sharding.update(
      mesh=manual_mesh,
      spec=core.modify_spec_for_auto_manual(aval.sharding.spec, manual_mesh))
  vma = (_spec_to_vma(spec) if check_vma else frozenset()) | aval.mat.varying
  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:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Set out_specs' unreduced equal to the output aval's unreduced (see aval.str_short(True))
2. Adjust the body so outputs are fully materialized (no unreduced) if out_specs must stay simple
3. Keep check_vma enabled during development to catch this early

Example fix

// before
out = shard_map(f, mesh, x, out_specs=P('d'))  # body yields unreduced 'r'
// after
out = shard_map(f, mesh, x, out_specs=P('d', unreduced=('r',)))

try: shard_map(...) except ValueError as e: if 'out_specs' in str(e) and 'unreduced' in str(e): adjust out_specs to match printed aval; else: raise

When it happens

Trigger: The function's outputs carry unreduced=('r',) semantics (e.g. from input mat info) but out_specs omits or changes unreduced.

Common situations: Writing manual/partial-manual shard_map bodies that return partially-replicated arrays while out_specs was written as a plain partition.

Related errors

• in_specs containing unreduced {spec} passed to shard_map sho
• 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/eb5529cf1a916b95. Report an issue: GitHub.