jax-ml/jax · error · ValueError

all_gather_reduced's input cannot be reduced across the axis

Error message

all_gather_reduced's input cannot be reduced across the axis_name provided. Got x={x_aval.str_short(True)} and {axis_name=}

What it means

`all_gather_reduced` requires its input to not already be reduced along the requested axis (`x.mat.reduced` must not intersect `axis_name`). An already-reduced value has already participated in a sum across that axis, so gathering it again would double-count semantics. The check in the abstract eval catches inconsistent collective pipelines early.

Source

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

  return tree_util.tree_map(bind, x)

all_gather_reduced_p = core.Primitive('all_gather_reduced')

def _all_gather_reduced_effectful_abstract_eval(
    x_aval, *, all_gather_dimension, axis_name, axis_size, tiled
):
  _check_axis_names(axis_name, 'all_gather_reduced')
  if not x_aval.mat.varying:
    raise ValueError('all_gather_reduced only accepts inputs that are'
                     f' varying. Got {x_aval.str_short(True)}')
  # If the intersection between x.mat.varying and axis_name is empty, error
  if not (x_aval.mat.varying & set(axis_name)):
    raise ValueError(
        'all_gather_reduced is a Varying -> Reduced collective. This means '
        f'that the {axis_name=} passed to `all_gather_reduced` must be present '
        f'in jax.typeof(x).mat.varying={x_aval.mat.varying}')
  if x_aval.mat.reduced & set(axis_name):
    raise ValueError(
        "all_gather_reduced's input cannot be reduced across the axis_name"
        f" provided. Got x={x_aval.str_short(True)} and {axis_name=}")

  new_shape = list(x_aval.shape)
  if tiled:
    new_shape[all_gather_dimension] *= axis_size
  else:
    new_shape.insert(all_gather_dimension, axis_size)

  if x_aval.mat.unreduced:
    check_unreduced_kind('all_gather_reduced', x_aval.mat, UnreducedKind.sum)
  new_reduced = x_aval.mat.reduced | frozenset(axis_name)
  out_vma = frozenset(v for v in x_aval.mat.varying if v not in axis_name)
  out_mat = x_aval.mat.update(varying=out_vma, reduced=new_reduced)
  return (x_aval.update(shape=new_shape, manual_axis_type=out_mat),
          {*map(core.NamedAxisEffect, axis_name)})
all_gather_reduced_p.def_effectful_abstract_eval(
    _all_gather_reduced_effectful_abstract_eval)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Restructure so the input is varying (not reduced) along axis_name, e.g. apply pvary to a fresh invariant value
Insert a barrier/reset by re-annotating with pvary on a non-reduced value
Use a different collective that accepts reduced inputs if that matches your intent

Example fix

// before
r = lax.psum(x, 'i')           # r.mat.reduced contains 'i'
y = lax.all_gather_reduced(r, axis_name='i', ...)
// after
r = lax.psum(x, 'i')
y = lax.all_gather_reduced(lax.pvary(r, 'i'), axis_name='i', ...)  # only if re-varying is intended

Defensive patterns

Strategy: validation

Validate before calling

if jax.typeof(x).mat.reduced & set(axis_name):
    x = lax.pvary(x, axis_name)  # re-annotate if re-gathering intended
y = lax.all_gather_reduced(x, axis_name=axis_name, ...)

Type guard

def not_reduced_on(x, axis_name) -> bool:
    return not (jax.typeof(x).mat.reduced & set(axis_name))

Try / catch

try:
    y = lax.all_gather_reduced(x, axis_name=axis_name, ...)
except ValueError as e:
    if 'cannot be reduced across' in str(e):
        raise RuntimeError(f'pipeline bug: {axis_name} already reduced') from e
    raise

Prevention

Track collective state per axis in code review
Unit-test collective pipelines with a tiny mesh to catch state conflicts

When it happens

Trigger: Feeding the output of a psum-like reduced collective (which marks the axis as reduced) directly into `all_gather_reduced` on the same axis name.

Common situations: Chaining collectives (psum then all_gather) while porting manual SPMD code; misunderstanding that reduced-ness is tracked in the type and persists through operations.

Related errors

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