jax-ml/jax · error · ValueError

unreduced_psum_scatter is a Unreduced -> Varying collective.

Error message

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

What it means

The `axis_name` passed to `unreduced_psum_scatter` must be present in `jax.typeof(x).mat.unreduced`; if the input is unreduced on other axes but not the requested one, JAX raises this error. The input's unreduced state and the scatter axis must agree. It usually means a mismatch between mesh axis names or a missing unreduced annotation.

Source

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

  axis_size = _axis_size(axis_name, None)
  def bind(leaf):
    return unreduced_reduce_scatter_p.bind(
        leaf, axis_name=axis_name, scatter_dimension=scatter_dimension,
        axis_size=axis_size, tiled=tiled)
  return tree_util.tree_map(bind, x)

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

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Print `jax.typeof(x).mat.unreduced` and pass one of those axes as axis_name
2. Add the missing unreduced producer for the intended axis
3. Verify axis names against the mesh/annotation definitions for typos

Example fix

// before
y = lax.unreduced_psum_scatter(x, axis_name='j', ...)
// after
y = lax.unreduced_psum_scatter(x, axis_name='i', ...)  # 'i' in x.mat.unreduced

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

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

try:
    y = lax.unreduced_psum_scatter(x, axis_name=axis_name, ...)
except ValueError as e:
    if 'Unreduced -> Varying' in str(e):
        axis_name = tuple(jax.typeof(x).mat.unreduced)[:1]
        y = lax.unreduced_psum_scatter(x, axis_name=axis_name, ...)
    else:
        raise

When it happens

Trigger: x is unreduced along 'i' but calling with `axis_name='j'`; or passing axis names that do not intersect `x.mat.unreduced`.

Common situations: Copy-pasting collective calls between model shards with different mesh axis names; typos in axis_name; renaming axes in the mesh without updating collective calls.

Related errors

• unbound axis name: {axis_name}
• all_gather_reduced is a Varying -> Reduced collective. This
• unreduced_psum_scatter only accepts inputs that are unreduce
• unreduced_psum_scatter's input cannot be varying across the
• {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/16f83fe59a0d5952. Report an issue: GitHub.