{"record":{"id":"bd50e762b4d43e50","repo":"jax-ml/jax","slug":"unreduced-in-prefix-specs-s-can-only-be-used-w","errorCode":null,"errorMessage":"unreduced in {prefix}_specs {s} can only be used when the mesh passed to shard_map contains axis names all of type `Explicit`. Got mesh {mesh}","messagePattern":"unreduced in (.+?)_specs (.+?) can only be used when the mesh passed to shard_map contains axis names all of type `Explicit`\\. Got mesh (.+?)","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/shard_map.py","lineNumber":462,"sourceCode":"\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\n  if error_type == SpecErrorType.input and specs is None:\n    raise TypeError(\n        \"shard_map in_specs argument must be a pytree of \"\n        \"`jax.sharding.PartitionSpec` instances, but it was None.\\n\"\n        \"Instead of `in_specs=None`, did you mean `in_specs=P()`, \"\n        \"where `P = jax.sharding.PartitionSpec`?\")","sourceCodeStart":444,"sourceCodeEnd":480,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/shard_map.py#L444-L480","documentation":"Raised when an in_specs/out_specs PartitionSpec uses the `unreduced` field while the mesh passed to shard_map has axis names that are not all of type `Explicit` (e.g. they are implicit/abstract axes). Unreduced specs require every referenced mesh axis to be explicit because replication semantics are only defined for explicit axes.","triggerScenarios":"Calling shard_map(mesh, in_specs=P('x', unreduced=('y',)), ...) where mesh axis 'y' is not an Explicit axis (e.g. created from an AbstractMesh or with non-explicit axis types).","commonSituations":"Using new-style abstract/implicit meshes (e.g. jax.experimental.mesh or context-managed meshes) together with the unreduced/reduced PartitionSpec extension; mixing API levels between JAX versions.","solutions":["Ensure every axis referenced in `unreduced` is an Explicit axis on the mesh","Construct the mesh with explicit axis names/sizes (jax.sharding.Mesh) instead of an abstract mesh","Remove `unreduced` from the spec if replication is not intended"],"exampleFix":"// before\nshard_map(f, mesh=abstract_mesh, in_specs=P('data', unreduced=('rep',)))\n// after\nmesh = jax.make_mesh((8,), ('rep',), axis_types=(AxisType.Explicit,))\nshard_map(f, mesh=mesh, in_specs=P('data', unreduced=('rep',)))","handlingStrategy":"validation","validationCode":"from jax.sharding import PartitionSpec as P\ndef axes_explicit(mesh, spec):\n    return all(mesh._name_to_type.get(u) == AxisType.Explicit for u in (spec.unreduced or ()))","typeGuard":"def has_only_explicit_unreduced(mesh, s) -> bool:\n    return all(u in mesh._name_to_type and mesh._name_to_type[u] == AxisType.Explicit for u in s.unreduced)","tryCatchPattern":"try: shard_map(...) except ValueError as e: if 'unreduced in' in str(e): rebuild mesh with explicit axes; else: raise","preventionTips":["Always create meshes with jax.make_mesh and explicit axis_types","Assert axis types before using unreduced/reduced specs"],"tags":["shard-map","partition-spec","mesh","jax","unreduced"],"backgroundTag":"invalid-sharding-spec","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}