jax-ml/jax · error · ValueError

unreduced in {prefix}_specs {s} can only be used when the me

Error message

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}

What it means

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.

Source

Thrown at jax/_src/shard_map.py:462

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
  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`?")

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Ensure every axis referenced in `unreduced` is an Explicit axis on the mesh
2. Construct the mesh with explicit axis names/sizes (jax.sharding.Mesh) instead of an abstract mesh
3. Remove `unreduced` from the spec if replication is not intended

Example fix

// before
shard_map(f, mesh=abstract_mesh, in_specs=P('data', unreduced=('rep',)))
// after
mesh = jax.make_mesh((8,), ('rep',), axis_types=(AxisType.Explicit,))
shard_map(f, mesh=mesh, in_specs=P('data', unreduced=('rep',)))

from jax.sharding import PartitionSpec as P
def axes_explicit(mesh, spec):
    return all(mesh._name_to_type.get(u) == AxisType.Explicit for u in (spec.unreduced or ()))

def has_only_explicit_unreduced(mesh, s) -> bool:
    return all(u in mesh._name_to_type and mesh._name_to_type[u] == AxisType.Explicit for u in s.unreduced)

try: shard_map(...) except ValueError as e: if 'unreduced in' in str(e): rebuild mesh with explicit axes; else: raise

When it happens

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

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

Related errors

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

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