{"record":{"id":"c2982b2e36c9cd25","repo":"jax-ml/jax","slug":"unreduced-reduced-can-only-be-passed-to-prefix-s","errorCode":null,"errorMessage":"unreduced/reduced can only be passed to {prefix}_specs when shard_map is in full manual mode. Got mesh axis names {mesh.axis_names}, manual_axes: {manual_axes}, specs: {s}. Please file a bug at https://github.com/jax-ml/jax/issues.","messagePattern":"unreduced/reduced can only be passed to (.+?)_specs when shard_map is in full manual mode\\. Got mesh axis names (.+?), manual_axes: (.+?), specs: (.+?)\\. Please file a bug at https://github\\.com/jax-ml/jax/issues\\.","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/shard_map.py","lineNumber":456,"sourceCode":"  return spec.update(partitions=tuple(out))\n\n\n# Error checking and messages\n\nSpecErrorType = enum.Enum('SpecErrorType', ['input', 'out'])\n\ndef _check_unreduced(error_type, mesh, manual_axes, specs):\n  from jax._src.hijax import HiPspec\n  prefix = 'in' if error_type == SpecErrorType.input else 'out'\n  full_manual = frozenset(mesh.axis_names) == manual_axes\n  specs_flat, _ = tree_flatten(specs)\n  for s in specs_flat:\n    if isinstance(s, HiPspec):\n      continue  # TODO(mattjj,yashkatariya): add user validation method\n    if not s.unreduced and not s.reduced:\n      continue\n    if not full_manual:\n      raise NotImplementedError(\n          f\"unreduced/reduced can only be passed to {prefix}_specs when\"\n          \" shard_map is in full manual mode. Got mesh axis names\"\n          f\" {mesh.axis_names}, manual_axes: {manual_axes}, specs: {s}. Please\"\n          \" file a bug at https://github.com/jax-ml/jax/issues.\")\n    if not all(mesh._name_to_type[u] == AxisType.Explicit for u in s.unreduced):\n      raise ValueError(\n          f\"unreduced in {prefix}_specs {s} can only be used when the mesh\"\n          \" passed to shard_map contains axis names all of type `Explicit`.\"\n          f\" Got mesh {mesh}\")\n    if not all(mesh._name_to_type[u] == AxisType.Explicit for u in s.reduced):\n      raise ValueError(\n          f\"reduced in {prefix}_specs {s} can only be used when the mesh\"\n          \" passed to shard_map contains axis names all of type `Explicit`.\"\n          f\" Got mesh {mesh}\")\n\n\ndef _check_specs(error_type: SpecErrorType, specs: Any, manual_axes) -> None:\n  from jax._src.hijax import HiPspec","sourceCodeStart":438,"sourceCodeEnd":474,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/shard_map.py#L438-L474","documentation":"The 'unreduced' and 'reduced' spec keywords (HiPspec extensions for controlling cross-shard reductions) are only legal when shard_map runs in full manual mode — i.e. manual_axes covering the whole mesh. Using them with an automatic/partial-manual shard_map raises NotImplementedError; the message invites filing a bug because the API surface is still experimental.","triggerScenarios":"Passing in_specs/out_specs entries containing Unreduced or Reduced (e.g. P('i', Reduced())) to shard_map without manual_axes set to all mesh axes.","commonSituations":"Experimenting with new SPMD reduction features from JAX docs/examples that assume full manual mode; copying pipeline-parallel or FSDP examples verbatim into code that uses automatic sharding; API drift across JAX versions as this feature evolves.","solutions":["Switch to full manual mode: pass manual_axes=frozenset(mesh.axis_names) (or all axes) to shard_map","Remove unreduced/reduced from the specs and perform reductions explicitly inside the function with jax.lax.psum / jax.lax.all_reduce"],"exampleFix":"// before\njax.shard_map(f, mesh=mesh, in_specs=P('i'), out_specs=P(Reduced('i')))(x)\n\n// after\njax.shard_map(f, mesh=mesh, manual_axes=frozenset(mesh.axis_names),\n              in_specs=P('i'), out_specs=P(Reduced('i')))(x)\n# or reduce manually:\njax.shard_map(lambda x: jax.lax.psum(x, 'i'), mesh=mesh,\n              in_specs=P('i'), out_specs=P())(x)","handlingStrategy":"validation","validationCode":"full_manual = manual_axes is not None and set(manual_axes) == set(mesh.axis_names)\nif uses_reduced_specs(in_specs) or uses_reduced_specs(out_specs):\n    assert full_manual, 'unreduced/reduced requires full manual mode'","typeGuard":null,"tryCatchPattern":"try:\n    out = shmapped(x)\nexcept NotImplementedError as e:\n    if 'full manual mode' in str(e):\n        shmapped = jax.shard_map(f, mesh=mesh, manual_axes=frozenset(mesh.axis_names),\n                                 in_specs=..., out_specs=...)\n        out = shmapped(x)\n    else: raise","preventionTips":["Only use unreduced/reduced together with manual_axes covering the whole mesh","Otherwise reduce explicitly with jax.lax.psum inside the function","Track JAX release notes; this API is experimental and shifting"],"tags":["jax","shard-map","manual-mode","unreduced","experimental"],"backgroundTag":"feature-requires-manual-mode","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}