jax-ml/jax · error · NotImplementedError

unreduced/reduced can only be passed to {prefix}_specs when

Error message

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.

What it means

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.

Source

Thrown at jax/_src/shard_map.py:456

  return spec.update(partitions=tuple(out))


# Error checking and messages

SpecErrorType = enum.Enum('SpecErrorType', ['input', 'out'])

def _check_unreduced(error_type, mesh, manual_axes, specs):
  from jax._src.hijax import HiPspec
  prefix = 'in' if error_type == SpecErrorType.input else 'out'
  full_manual = frozenset(mesh.axis_names) == manual_axes
  specs_flat, _ = tree_flatten(specs)
  for s in specs_flat:
    if isinstance(s, HiPspec):
      continue  # TODO(mattjj,yashkatariya): add user validation method
    if not s.unreduced and not s.reduced:
      continue
    if not full_manual:
      raise NotImplementedError(
          f"unreduced/reduced can only be passed to {prefix}_specs when"
          " shard_map is in full manual mode. Got mesh axis names"
          f" {mesh.axis_names}, manual_axes: {manual_axes}, specs: {s}. Please"
          " file a bug at https://github.com/jax-ml/jax/issues.")
    if not all(mesh._name_to_type[u] == AxisType.Explicit for u in s.unreduced):
      raise ValueError(
          f"unreduced in {prefix}_specs {s} can only be used when the mesh"
          " passed to shard_map contains axis names all of type `Explicit`."
          f" Got mesh {mesh}")
    if not all(mesh._name_to_type[u] == AxisType.Explicit for u in s.reduced):
      raise ValueError(
          f"reduced in {prefix}_specs {s} can only be used when the mesh"
          " passed to shard_map contains axis names all of type `Explicit`."
          f" Got mesh {mesh}")


def _check_specs(error_type: SpecErrorType, specs: Any, manual_axes) -> None:
  from jax._src.hijax import HiPspec

View on GitHub (pinned to 1e1c6a8fc0)

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

Example fix

// before
jax.shard_map(f, mesh=mesh, in_specs=P('i'), out_specs=P(Reduced('i')))(x)

// after
jax.shard_map(f, mesh=mesh, manual_axes=frozenset(mesh.axis_names),
              in_specs=P('i'), out_specs=P(Reduced('i')))(x)
# or reduce manually:
jax.shard_map(lambda x: jax.lax.psum(x, 'i'), mesh=mesh,
              in_specs=P('i'), out_specs=P())(x)

Defensive patterns

Strategy: validation

Validate before calling

full_manual = manual_axes is not None and set(manual_axes) == set(mesh.axis_names)
if uses_reduced_specs(in_specs) or uses_reduced_specs(out_specs):
    assert full_manual, 'unreduced/reduced requires full manual mode'

Try / catch

try:
    out = shmapped(x)
except NotImplementedError as e:
    if 'full manual mode' in str(e):
        shmapped = jax.shard_map(f, mesh=mesh, manual_axes=frozenset(mesh.axis_names),
                                 in_specs=..., out_specs=...)
        out = shmapped(x)
    else: raise

Prevention

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

When it happens

Trigger: 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.

Common situations: 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.

Related errors

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