jax-ml/jax · error · ValueError

{name} only accepts inputs that are unreduced. Got {aval.str

Error message

{name} only accepts inputs that are unreduced. Got {aval.str_short(True)}

What it means

JAX's unreduced collectives (`unreduced_psum`, `unreduced_pmax`, `unreduced_pmin`) require their input to actually be marked as 'unreduced' along the corresponding mesh axes in the sharding annotation (the `mat.unreduced` set of the aval). This ValueError fires when the abstract-eval sees an input whose unreduced set is empty, i.e. the value was never produced/constrained to be unreduced.

Source

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

                       partial(_unreduced_reduce_scatter_lowering, lax.add_p))

############################## unreduced_psum ###########################

# 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)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Ensure the input was produced by a pipeline that marks it unreduced (e.g. the output of an operation returning unreduced values along that axis)
2. Use the standard collective (`psum`, `pmax`, `pmin`) instead of the unreduced variant if your input is invariant
3. Check that the axis_name matches the mesh axis along which the value is actually unreduced
4. Inspect `jax.typeof(x).mat` before the call to confirm the unreduced set

Example fix

// before
y = jax.lax.unreduced_psum(x, 'data')  # x is invariant
// after
y = jax.lax.psum(x, 'data')  # x is not unreduced; use standard psum

import jax
t = jax.typeof(x)
assert t.mat.unreduced, f'input not unreduced: {t.mat}'

def is_unreduced(x) -> bool:
    return bool(jax.typeof(x).mat.unreduced)

When it happens

Trigger: Calling `jax.lax.psum(x, axis_name)` on a plain (fully materialized / invariant) array under a named axis context where the API routes to the unreduced variant; passing a value to `unreduced_psum`/`unreduced_pmax`/`unreduced_pmin` that has `mat.unreduced` empty.

Common situations: Mixing the newer unreduced-collective API with code that assumes values are replicated/invariant; sharding-annotation mismatch after upgrading JAX to a version with mat (materialization) tracking; forgetting to produce the input via an operation that marks it unreduced.

Related errors

• Axes mentioned in `manual_axis_type` field of ShapedArray sh
• unbound axis name: {axis_name}
• {name} is a Unreduced -> Invariant collective. This means th
• {name}'s input cannot be varying across the axis_name provi
• preduced is a Invariant->Reduced collective. This means that

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