jax-ml/jax · error · ValueError

Reduced axes can only refer to mesh axes that is of type `Ex

Error message

Reduced axes can only refer to mesh axes that is of type `Explicit`. Got reduced axes: {pspec.reduced} and mesh: {mesh}

What it means

Raised by jax's NamedSharding PSpec validation (check_pspec/_check_mesh_unreduced) when a PartitionSpec marks axes as 'reduced', but the corresponding mesh axis is not an Explicit axis (it is Auto or Manual, i.e. created via automatic or manual/spmd partitioning rather than explicitly named in the Mesh). Reduction semantics are only defined over explicitly-declared mesh axes.

Source

Thrown at jax/_src/named_sharding.py:605

def _check_mesh_unreduced(mesh, pspec):
  for u in pspec.unreduced:
    if u not in mesh.axis_names:
      raise ValueError(
          f'Unreduced axes {u} is not found in {mesh.axis_names=}. '
          f'Got {pspec=}')
    if mesh._name_to_type[u] in (AxisType.Auto, AxisType.Manual):
      raise ValueError(
          'Unreduced axes can only refer to mesh axes that is of type'
          f' `Explicit`. Got unreduced axes: {pspec.unreduced} and'
          f' mesh: {mesh}')

  for u in pspec.reduced:
    if u not in mesh.axis_names:
      raise ValueError(
          f'Reduced axes {u} is not found in {mesh.axis_names=}. '
          f'Got {pspec=}')
    if mesh._name_to_type[u] in (AxisType.Auto, AxisType.Manual):
      raise ValueError(
          'Reduced axes can only refer to mesh axes that is of type'
          f' `Explicit`. Got reduced axes: {pspec.reduced} and'
          f' mesh: {mesh}')

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Use a Mesh whose axes are all explicit (jax.sharding.Mesh with named device lists) so the reduced axis resolves to AxisType.Explicit
2. Remove the reduced axes from the PSpec and do the reduction manually (e.g. jax.lax.psum / .sum) instead of relying on unreduced/reduced spec semantics
3. If on Auto axes, rely on the GSPMD auto partitioner instead of NamedSharding reduced specs

Example fix

// before
pspec = P(None, 'x', reduced=('x',))  # mesh axis 'x' is Auto/Manual
sharding = NamedSharding(mesh, pspec)

// after
out = jax.pmap(lambda x: x.sum('x'))(arr)  # explicit reduction
# or use an all-Explicit Mesh:
mesh = jax.sharding.Mesh(jax.devices(), ('x',))
sharding = NamedSharding(mesh, P(None, 'x', reduced=('x',)))

from jax.experimental.mesh_utils import ...  # if needed
for u in pspec.reduced:
    assert u in mesh.axis_names, f'{u} not in mesh'
    # mirror the internal check without private access: ensure the axis was
    # created explicitly via jax.sharding.Mesh(devices, (names...))
assert all(a in mesh.axis_names for a in pspec.reduced)

When it happens

Trigger: Passing a NamedSharding whose PSpec uses reduced axes (e.g. P(None, 'x', reduced=('x',))) where mesh axis 'x' was created as AxisType.Auto (from jax.experimental.mesh_utils or auto-partitioned meshes) or AxisType.Manual (from spmd manual partitioning), then using it as an out_sharding.

Common situations: Migrating SPMD/manual-partition code to NamedSharding out_sharding APIs; mixing Mesh created with auto axes (e.g. openxla mesh_utils) with PSpec reduced syntax; specifying out_sharding on initializers/one_hot with a mesh built for auto-sharding.

Related errors

• Mapped away dimension of inputs passed to vmap should be sha
• Sharding spec {spec} implies that array axis {dim} is partit
• The error code state and the predicate must be on the same m
• out_sharding passed to {name} cannot be reduced. Got {out_sh
• Unmapped values passed to vmap cannot be sharded along the m

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