jax-ml/jax · error · ValueError

all_gather_reduced only accepts inputs that are varying. Got

Error message

all_gather_reduced only accepts inputs that are varying. Got {x_aval.str_short(True)}

What it means

`all_gather_reduced` gathers an array that is varying (replicated-but-named) along the given axis, so its input must have a non-empty `mat.varying` set. The abstract eval rejects inputs with no varying axes because there is nothing to gather along. This almost always means the producer of the input never marked it varying (missing `pvary`).

Source

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

    return x
  axis_size = _axis_size(axis_name, None)
  def bind(leaf):
    prim = all_gather_reduced_start_p if is_async else all_gather_reduced_p
    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)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Apply `lax.pvary(x, 'i')` before `all_gather_reduced`
2. Verify the producer of x actually varies the axis (check `jax.typeof(x).mat.varying`)
3. Ensure you meant all_gather_reduced rather than a plain all_gather on a sharded array

Example fix

// before
y = lax.all_gather_reduced(x, axis_name='i', all_gather_dimension=0, axis_size=8, tiled=False)
// after
x = lax.pvary(x, 'i')
y = lax.all_gather_reduced(x, axis_name='i', all_gather_dimension=0, axis_size=8, tiled=False)

if not jax.typeof(x).mat.varying:
    x = lax.pvary(x, axis_name)
y = lax.all_gather_reduced(x, axis_name=axis_name, ...)

def is_varying(x) -> bool:
    return bool(jax.typeof(x).mat.varying)

try:
    y = lax.all_gather_reduced(x, axis_name=axis_name, ...)
except ValueError as e:
    if 'only accepts inputs that are varying' in str(e):
        y = lax.all_gather_reduced(lax.pvary(x, axis_name), axis_name=axis_name, ...)
    else:
        raise

When it happens

Trigger: Calling `lax.all_gather_reduced(x, axis_name='i', ...)` on an input whose abstract value has `mat.varying == frozenset()`, e.g. a plain array or output of a fully reduced collective without a preceding pvary.

Common situations: Migrating psum/pmap code to named-axis SPMD style and forgetting the `pvary` annotation; reordering collectives so the varying-producing op is dropped or optimized away.

Related errors

• pvary is a invariant->variant collective. This means that th
• all_gather_reduced is a Varying -> Reduced collective. This
• all_gather_reduced's input cannot be reduced across the axis
• unreduced_psum_scatter only accepts inputs that are unreduce
• callbacks are only supported in spmd computations when all m

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