jax-ml/jax · error · ValueError

shard_map out_specs rank error (msg from _spec_rank_error)

Error message

shard_map out_specs rank error (msg from _spec_rank_error)

What it means

After the wrapped function in shard_map runs, out_specs is used to reassemble per-shard outputs into a global array; each shard's rank must equal the number of mesh axes named in its PartitionSpec. This ValueError (message built by _spec_rank_error) means an output shard has too few or too many dimensions for its spec — including the special case of rank-0 outputs that vary across the mesh, which the appended hint addresses.

Source

Thrown at jax/_src/shard_map.py:349

        val = pvary(val, tuple(_spec_to_vma(spec) - aval.mat.varying))
        return val
      if check_vma:
        ans_ft = ans_ft.map2(out_specs_flat, add_implicit_pvary)
      return ans_ft.with_aux(out_specs_flat)

    try:
      newly_manual_axes = axis_names - set(mesh.manual_axes)
      out_ft = shard_map_p.bind(
          *dyn_args, subfuns=(f_wrapped,), mesh=mesh, in_specs=in_specs_flat,
          check_vma=check_vma, newly_manual_axes=newly_manual_axes, debug_info=dbg)
    except _SpecError as e:
      fails, out_tree = e.args
      msg = _spec_rank_error(SpecErrorType.out, f, out_tree, out_specs, fails)
      if any(fail is not no_fail and not fail.shape for fail in fails):
        msg += (" In particular, for rank 0 outputs which are not constant "
                "over the mesh, add at least one (singleton) axis to them so "
                "that they can be concatenated using out_specs.")
      raise ValueError(msg) from None
    except _RepError as e:
      fails, out_tree, = e.args
      msg = _inout_vma_error(f, mesh, out_tree, out_specs, fails)
      raise ValueError(msg) from None
    return out_ft.unflatten()
  return cast(F, wrapped)


def _axes_to_pspec(axis_name, axis):
  if axis is None:
    return P()
  return P(*[None] * axis + [axis_name])


def _shmap_checks(mesh, axis_names, in_specs, out_specs, _smap):
  if mesh is None:
    mesh = get_abstract_mesh()
    if mesh.empty:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Make the function return shards whose rank equals the count of named axes in the corresponding out_specs entry (add singleton dims if needed: out[None, None])
2. Adjust out_specs to match actual output rank (e.g. P('i') instead of P('i','j'))
3. For a global scalar that IS constant over the mesh, return jnp.asarray(scalar) with no axes in the spec, or use P() / axis-reduction via 'reduced' specs in manual mode

Example fix

// before
def f(x): return x.sum()  # rank-0 shard
jax.shard_map(f, mesh=mesh, in_specs=P('i'), out_specs=P('i'))(x)

// after
def f(x): return x.sum(keepdims=True)  # rank-1 shard
jax.shard_map(f, mesh=mesh, in_specs=P('i'), out_specs=P('i'))(x)

# each shard output's ndim must equal len of its out_specs entry
spec_entry = P('i', 'j')
def check_shard_rank(out_shard, spec_entry):
    assert out_shard.ndim == len([a for a in spec_entry if a is not None]), (
        out_shard.ndim, spec_entry)

try:
    out = shmapped(x)
except ValueError as e:
    if 'out_specs rank error' in str(e):
        # add singleton dims to outputs or simplify out_specs, then retry
        ...
    raise

When it happens

Trigger: A mapped function returning a scalar or per-shard tensor whose ndim doesn't match len(out_specs entry), e.g. out_specs=P('i','j') with the function returning rank-1 shards, or returning a non-constant Python/0-d value on a multi-axis mesh.

Common situations: Returning loss scalars or per-shard statistics from inside shard_map; changing the function's output shape without updating out_specs; using 'unreduced' semantics where a manual axis is expected.

Related errors

• The process id of the current process must be defined.
• process_id must be a nonnegative int. Got process_id={proces
• smap axis_name should be a `str` or a `Hashable`, but got {a
• smap in_axes must be an int, None, jax.sharding.Infer, or a
• in_specs passed to shard_map: {s} does not match the specs o

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