jax-ml/jax · error · ValueError

all_gather_reduced is a Varying -> Reduced collective. This

Error message

all_gather_reduced is a Varying -> Reduced collective. This means that the {axis_name=} passed to `all_gather_reduced` must be present in jax.typeof(x).mat.varying={x_aval.mat.varying}

What it means

`all_gather_reduced` is a Varying -> Reduced collective: the `axis_name` you pass must appear in `jax.typeof(x).mat.varying`. If the input varies along other axes but not the requested one, gathering along it is meaningless, so the abstract eval raises this error. It indicates the array and the collective disagree about which mesh axis to operate on.

Source

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

    return prim.bind(
        leaf,
        all_gather_dimension=canonicalize_axis(
            axis, np.ndim(leaf) if tiled else np.ndim(leaf) + 1),
        axis_name=axis_name, axis_size=axis_size, tiled=tiled)
  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)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Make axis_name match an axis in `jax.typeof(x).mat.varying` (print it to confirm)
2. Add `lax.pvary(x, axis_name)` if the axis was never varied
3. Check for axis-name typos against the mesh definition

Example fix

// before
x = lax.pvary(x, 'i')
y = lax.all_gather_reduced(x, axis_name='j', ...)
// after
x = lax.pvary(x, 'i')
y = lax.all_gather_reduced(x, axis_name='i', ...)

assert set(jax.typeof(x).mat.varying) & set(axis_name), (
    f'axis_name {axis_name} not in {jax.typeof(x).mat.varying}')
y = lax.all_gather_reduced(x, axis_name=axis_name, ...)

def axis_is_varying(x, axis_name) -> bool:
    return bool(set(jax.typeof(x).mat.varying) & set(axis_name))

try:
    y = lax.all_gather_reduced(x, axis_name=axis_name, ...)
except ValueError as e:
    if 'Varying -> Reduced' in str(e):
        correct = next(iter(jax.typeof(x).mat.varying))
        y = lax.all_gather_reduced(x, axis_name=correct, ...)
    else:
        raise

When it happens

Trigger: Input varied along axis 'i' but calling `all_gather_reduced(x, axis_name='j', ...)`; passing a tuple of axis names none of which intersect `x.mat.varying`.

Common situations: Renaming mesh axes or copy-pasting collective calls between pipelines with different axis names; typos in axis_name strings.

Related errors

• unbound axis name: {axis_name}
• all_gather_reduced only accepts inputs that are varying. Got
• all_gather_reduced's input cannot be reduced across the axis
• unreduced_psum_scatter is a Unreduced -> Varying collective.
• {name} is a Unreduced -> Invariant collective. This means th

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