jax-ml/jax · error · ValueError

unreduced_psum_scatter only accepts inputs that are unreduce

Error message

unreduced_psum_scatter only accepts inputs that are unreduced. Got {x_aval.str_short(True)}

What it means

`unreduced_psum_scatter` is an Unreduced -> Varying collective: it consumes a value whose `mat.unreduced` set is non-empty (e.g. the output of a psum-like computation tracked as unreduced-sum). If the input has no unreduced axes, the abstract eval rejects it immediately. Typically the producer op that marks the value unreduced is missing.

Source

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

  if not isinstance(axis_name, tuple):
    axis_name = (axis_name,)
  if not axis_name:
    return x
  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:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Ensure the input comes from a computation marked unreduced (e.g. proper collective output) on the target axis
2. Check `jax.typeof(x).mat.unreduced` before the call and add the missing producer op
3. Use regular `psum_scatter` if you do not need unreduced semantics

Example fix

// before
y = lax.unreduced_psum_scatter(x, axis_name='i', scatter_dimension=0, axis_size=8)
// after
# produce x so that jax.typeof(x).mat.unreduced contains 'i', e.g. via an unreduced-sum collective
y = lax.unreduced_psum_scatter(unreduced_x, axis_name='i', scatter_dimension=0, axis_size=8)

if not jax.typeof(x).mat.unreduced:
    raise ValueError(f'input not unreduced: {jax.typeof(x)}')
y = lax.unreduced_psum_scatter(x, ...)

def is_unreduced(x) -> bool:
    return bool(jax.typeof(x).mat.unreduced)

try:
    y = lax.unreduced_psum_scatter(x, ...)
except ValueError as e:
    if 'only accepts inputs that are unreduced' in str(e):
        y = lax.psum_scatter(x, axis_name)  # fallback to standard collective
    else:
        raise

When it happens

Trigger: Passing a plain array or a fully varying value into `unreduced_psum_scatter` without a preceding op producing `UnreducedKind.sum` state on the relevant axis.

Common situations: Building reduce-scatter pipelines manually in the new named-mesh API and skipping the unreduced-producing step; refactoring from `psum_scatter` where the unreduced tracking was implicit.

Related errors

• pvary is a invariant->variant collective. This means that th
• all_gather_reduced only accepts inputs that are varying. Got
• unreduced_psum_scatter is a Unreduced -> Varying collective.
• unreduced_psum_scatter's input cannot be varying across the
• 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/70db65f4539bced2. Report an issue: GitHub.