jax-ml/jax · error · ValueError

unreduced_psum_scatter's input cannot be varying across the

Error message

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

What it means

`unreduced_psum_scatter`'s input must not already be varying along the requested `axis_name` (i.e. `x.mat.varying` must not intersect it). A value cannot be simultaneously varying and about to become varying via scatter on the same axis. The abstract eval rejects this inconsistent state.

Source

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

unreduced_reduce_scatter_p = core.Primitive('unreduced_reduce_scatter')

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

  new_shape = list(x_aval.shape)
  scatter_dim_input_size = x_aval.shape[scatter_dimension]
  if tiled:
    if scatter_dim_input_size % axis_size != 0:
      raise ValueError(f"tiled reduce_scatter operand scatter dimension size "
                       f"{scatter_dim_input_size} must be divisible by "
                       f"shard_count {axis_size}")
    new_shape[scatter_dimension] = scatter_dim_input_size // axis_size
  else:
    if scatter_dim_input_size != axis_size:
      raise ValueError(f"reduce_scatter operand scatter dimension size "
                       f"{scatter_dim_input_size} must match shard count "
                       f"{axis_size}")
    del new_shape[scatter_dimension]

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Remove the earlier pvary on that axis, or scatter along a different axis not in `x.mat.varying`
2. Restructure so the value is unreduced (not varying) on the scatter axis before the call
3. Inspect `jax.typeof(x).mat` to see the full varying/unreduced state and fix the producer chain

Example fix

// before
x = lax.pvary(x, 'i')
y = lax.unreduced_psum_scatter(x, axis_name='i', ...)
// after
y = lax.unreduced_psum_scatter(unreduced_x, axis_name='i', ...)  # x not varying on 'i'

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

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

try:
    y = lax.unreduced_psum_scatter(x, axis_name=axis_name, ...)
except ValueError as e:
    if 'cannot be varying across' in str(e):
        raise RuntimeError(f'state conflict on {axis_name}; fix producer chain') from e
    raise

When it happens

Trigger: Input previously marked with pvary along 'i' then passed to `unreduced_psum_scatter(..., axis_name='i')`.

Common situations: Composing pvary with reduce-scatter pipelines where the same axis ends up in both varying and unreduced tracks; porting older sharding code that applied explicit sharding annotations before collective calls.

Related errors

• unreduced_psum_scatter only accepts inputs that are unreduce
• unreduced_psum_scatter is a Unreduced -> Varying collective.
• {name}'s input cannot be varying across the axis_name provi
• unbound axis name: {axis_name}
• pvary is a invariant->variant collective. This means that th

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