jax-ml/jax · error · ValueError

{name} is a Unreduced -> Invariant collective. This means th

Error message

{name} is a Unreduced -> Invariant collective. This means that the {axes=} passed to `{name}` must be present in jax.typeof(x).mat.unreduced={aval.mat.unreduced}

What it means

`unreduced_psum`/`unreduced_pmax`/`unreduced_pmin` are Unreduced -> Invariant collectives: the axis names passed to them must already appear in the input's `mat.unreduced` set. This ValueError fires when the intersection of `aval.mat.unreduced` and the requested axes is empty — the value is unreduced along some axes, but not the ones you asked to reduce.

Source

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

# Unreduced -> Invariant collective
def unreduced_psum(x, axis_name):
  if not isinstance(axis_name, (tuple, list)):
    axis_name = (axis_name,)
  if not axis_name:
    return x
  return tree_util.tree_map(
      lambda leaf: unreduced_psum_p.bind(leaf, axes=tuple(axis_name)), x)

unreduced_psum_p = core.Primitive('unreduced_psum')

def _unreduced_psum_pmax_pmin_abstract_eval(name, out_u_kind, aval, *, axes):
  _check_axis_names(axes, name)
  if not aval.mat.unreduced:
    raise ValueError(f'{name} only accepts inputs that are'
                     f' unreduced. Got {aval.str_short(True)}')
  # If intersection between x.unreduced & axis_name is empty, error
  if not (aval.mat.unreduced & frozenset(axes)):
    raise ValueError(
        f"{name} is a Unreduced -> Invariant collective. This"
        f" means that the {axes=} passed to `{name}` must"
        " be present in"
        f" jax.typeof(x).mat.unreduced={aval.mat.unreduced}")
  if aval.mat.varying & set(axes):
    raise ValueError(
        f"{name}'s input cannot be varying across the "
        f" axis_name provided. Got x={aval.str_short(True)} and {axes=}")

  if any(isinstance(a, int) for a in axes):
    raise ValueError(f'{name} does not accept integer axis_name.'
                     f' Got axis_name={axes}')

  core.check_avals_context_mesh([aval], name)
  check_unreduced_kind(name, aval.mat, out_u_kind)
  out_u = frozenset(u for u in aval.mat.unreduced if u not in axes)
  kind = aval.mat.unreduced_kind if out_u else None
  out_mat = aval.mat.update(unreduced=out_u, unreduced_kind=kind)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Print `jax.typeof(x).mat.unreduced` and pass one of those axis names to the collective
2. Fix the axis name typo / use the correct mesh axis the value is unreduced along
3. Adjust upstream code so the value is produced unreduced along the intended axis

Example fix

// before
jax.lax.unreduced_psum(x, 'data')  # x unreduced along 'batch'
// after
jax.lax.unreduced_psum(x, 'batch')

import jax
axes = {'data'}
assert axes & set(jax.typeof(x).mat.unreduced), 'axis not in x.mat.unreduced'

def axis_is_unreduced(x, axis: str) -> bool:
    return axis in jax.typeof(x).mat.unreduced

When it happens

Trigger: Calling `unreduced_psum(x, 'data')` when `jax.typeof(x).mat.unreduced` is e.g. `{'batch'}` but not `'data'`; passing an axis_name not present in the input's unreduced set (even though the set itself is non-empty).

Common situations: Typos or renames of mesh axis names between model definition and collective call; using the wrong axis from a multi-axis mesh; refactors that change which axis a tensor is unreduced along.

Related errors

• unbound axis name: {axis_name}
• all_gather_reduced is a Varying -> Reduced collective. This
• unreduced_psum_scatter is a Unreduced -> Varying collective.
• {name} only accepts inputs that are unreduced. Got {aval.str
• {name}'s input cannot be varying across the axis_name provi

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