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.

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'
Defensive patterns

Strategy: validation

Validate before calling

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, ...)

Type guard

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

Try / catch

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

Prevention

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

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

Appendix: 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)