{"record":{"id":"edcef3328ac613cb","repo":"jax-ml/jax","slug":"shard-map-out-specs-rank-error-msg-from-spec-ran","errorCode":null,"errorMessage":"shard_map out_specs rank error (msg from _spec_rank_error)","messagePattern":"shard_map out_specs rank error \\(msg from _spec_rank_error\\)","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/shard_map.py","lineNumber":349,"sourceCode":"        val = pvary(val, tuple(_spec_to_vma(spec) - aval.mat.varying))\n        return val\n      if check_vma:\n        ans_ft = ans_ft.map2(out_specs_flat, add_implicit_pvary)\n      return ans_ft.with_aux(out_specs_flat)\n\n    try:\n      newly_manual_axes = axis_names - set(mesh.manual_axes)\n      out_ft = shard_map_p.bind(\n          *dyn_args, subfuns=(f_wrapped,), mesh=mesh, in_specs=in_specs_flat,\n          check_vma=check_vma, newly_manual_axes=newly_manual_axes, debug_info=dbg)\n    except _SpecError as e:\n      fails, out_tree = e.args\n      msg = _spec_rank_error(SpecErrorType.out, f, out_tree, out_specs, fails)\n      if any(fail is not no_fail and not fail.shape for fail in fails):\n        msg += (\" In particular, for rank 0 outputs which are not constant \"\n                \"over the mesh, add at least one (singleton) axis to them so \"\n                \"that they can be concatenated using out_specs.\")\n      raise ValueError(msg) from None\n    except _RepError as e:\n      fails, out_tree, = e.args\n      msg = _inout_vma_error(f, mesh, out_tree, out_specs, fails)\n      raise ValueError(msg) from None\n    return out_ft.unflatten()\n  return cast(F, wrapped)\n\n\ndef _axes_to_pspec(axis_name, axis):\n  if axis is None:\n    return P()\n  return P(*[None] * axis + [axis_name])\n\n\ndef _shmap_checks(mesh, axis_names, in_specs, out_specs, _smap):\n  if mesh is None:\n    mesh = get_abstract_mesh()\n    if mesh.empty:","sourceCodeStart":331,"sourceCodeEnd":367,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/shard_map.py#L331-L367","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","solutions":["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])","Adjust out_specs to match actual output rank (e.g. P('i') instead of P('i','j'))","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"],"exampleFix":"// before\ndef f(x): return x.sum()  # rank-0 shard\njax.shard_map(f, mesh=mesh, in_specs=P('i'), out_specs=P('i'))(x)\n\n// after\ndef f(x): return x.sum(keepdims=True)  # rank-1 shard\njax.shard_map(f, mesh=mesh, in_specs=P('i'), out_specs=P('i'))(x)","handlingStrategy":"validation","validationCode":"# each shard output's ndim must equal len of its out_specs entry\nspec_entry = P('i', 'j')\ndef check_shard_rank(out_shard, spec_entry):\n    assert out_shard.ndim == len([a for a in spec_entry if a is not None]), (\n        out_shard.ndim, spec_entry)","typeGuard":null,"tryCatchPattern":"try:\n    out = shmapped(x)\nexcept ValueError as e:\n    if 'out_specs rank error' in str(e):\n        # add singleton dims to outputs or simplify out_specs, then retry\n        ...\n    raise","preventionTips":["Keep returned shard ranks equal to named-axis counts in out_specs","Use keepdims=True when reducing inside shard_map","For constants over the mesh return P()-compatible scalars"],"tags":["jax","shard-map","out-specs","rank","distributed"],"backgroundTag":"sharding-spec-rank-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}