jax-ml/jax · error · ValueError

reduced in {prefix}_specs {s} can only be used when the mesh

Error message

reduced 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}

What it means

Raised when an in_specs/out_specs PartitionSpec uses the `reduced` field while the mesh passed to shard_map contains axis names that are not all of type `Explicit`. Reduced (collectively-summed) axes only have well-defined semantics on explicit mesh axes.

Source

Thrown at jax/_src/shard_map.py:467

  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
  if error_type == SpecErrorType.input and specs is None:
    raise TypeError(
        "shard_map in_specs argument must be a pytree of "
        "`jax.sharding.PartitionSpec` instances, but it was None.\n"
        "Instead of `in_specs=None`, did you mean `in_specs=P()`, "
        "where `P = jax.sharding.PartitionSpec`?")

  def check_spec(p):
    if isinstance(p, HiPspec):
      return True  # TODO(mattjj,yashkatariya): add user validation method
    if not isinstance(p, PartitionSpec):

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Make all axes referenced in `reduced` Explicit on the mesh
2. Use a concrete jax.sharding.Mesh with explicit axis types
3. Drop `reduced` from the spec and handle reduction manually inside the function

Example fix

// before
shard_map(f, mesh=amesh, in_specs=P(reduced=('y',)))
// after
mesh = jax.make_mesh((4,), ('y',), axis_types=(AxisType.Explicit,))
shard_map(f, mesh=mesh, in_specs=P(reduced=('y',)))

def axes_explicit(mesh, spec):
    return all(mesh._name_to_type.get(r) == AxisType.Explicit for r in (spec.reduced or ()))

def has_only_explicit_reduced(mesh, s) -> bool:
    return all(r in mesh._name_to_type and mesh._name_to_type[r] == AxisType.Explicit for r in s.reduced)

try: shard_map(...) except ValueError as e: if 'reduced in' in str(e): make axes explicit; else: raise

When it happens

Trigger: Calling shard_map with a spec like P('x', reduced=('y',)) where mesh axis 'y' is implicit/non-explicit (abstract mesh or non-explicit axis type).

Common situations: Combining the reduced/unreduced spec extension with abstract meshes or partially non-explicit axis types introduced in newer JAX versions.

Related errors

• unreduced in {prefix}_specs {s} can only be used when the me
• Mesh must be provided for shard_map with checkify.
• Found an unbound axis name: {name}. To fix this, please call
• all_gather_reduced's input cannot be reduced across the axis
• When `check_vma=True` on `jax.shard_map`, `manual_axis_type`

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