jax-ml/jax · error · ValueError

{name} cannot accept args with unreduced_kind={mat.unreduced

Error message

{name} cannot accept args with unreduced_kind={mat.unreduced_kind}. Expected unreduced_kind={kind}

What it means

jax.lax parallel collectives operate on arrays that may be 'unreduced' or 'reduced' with respect to a mapped axis (used by shard_map / SPMD typed collectives). check_unreduced_kind guards each primitive so it only receives the kind of array it can semantically handle — e.g. all_gather's reduced-effectful path expects unreduced inputs. A mismatch means the collective was applied to an array whose reduction status under that axis does not fit the operation.

Source

Thrown at jax/_src/lax/parallel.py:64

from jax._src.lax import lax
from jax._src.lax import slicing
from jax._src.lib.mlir import ir
from jax._src.lib.mlir.dialects import hlo
from jax._src.typing import Array
from jax._src.util import (canonicalize_axis, moveaxis, safe_map, safe_zip,
                           unzip2)
from jax._src.lib.mlir.dialects import func as func_dialect
import numpy as np

unsafe_map, map = map, safe_map
unsafe_zip, zip = zip, safe_zip


### parallel traceables

def check_unreduced_kind(name, mat, kind):
  if mat.unreduced_kind is not kind:
    raise ValueError(
        f"{name} cannot accept args with unreduced_kind={mat.unreduced_kind}."
        f" Expected unreduced_kind={kind}")

def psum(x, axis_name, *, axis_index_groups=None):
  """Compute an all-reduce sum on ``x`` over the pmapped axis ``axis_name``.

  If ``x`` is a pytree then the result is equivalent to mapping this function to
  each leaf in the tree.

  Inputs of boolean dtype are converted to integers before the reduction.

  Args:
    x: array(s) with a mapped axis named ``axis_name``.
    axis_name: hashable Python object used to name a pmapped axis (see the
      :func:`jax.pmap` documentation for more details).
    axis_index_groups: optional list of lists containing axis indices (e.g. for
      an axis of size 4, [[0, 1], [2, 3]] would perform psums over the first
      two and last two replicas). Groups must cover all axis indices exactly

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Re-check the order of collectives in your shard_map body; ensure the input to each collective has the expected reduction state (e.g. don't feed all_gather-reduced output back into the unreduced path)
Use plain lax.psum / all_gather traceables instead of lower-level typed primitives so JAX tracks state for you
Upgrade/downgrade to a JAX version whose shard_map semantics match the code (these rules changed across releases)

Example fix

// before
y = all_gather_reduced(psum_out, axis_name)  # psum output is reduced -> mismatch

// after
y = jax.lax.all_gather(x, axis_name)  # operate on the original unreduced x

Defensive patterns

Strategy: validation

Validate before calling

# prefer public traceables (jax.lax.psum etc.) so unreduced_kind is tracked automatically

Prevention

Avoid low-level typed-collective primitives directly
Keep collective order simple: one reduce per axis per shard_map body

When it happens

Trigger: Mixing typed SPMD collectives, e.g. calling the reduced all_gather / psum / pmax path on an array already marked unreduced (or vice versa) for axis_name — typically inside shard_map pipelines composing psum, all_gather, and reduce_scatter in one function.

Common situations: Refactoring shard_map code where a previous collective changed the array's unreduced_kind; upgrading JAX versions where typed-collective rules tightened; hand-writing compose pipelines of psum -> all_gather -> reduce_scatter.

Related errors

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