jax-ml/jax · error · ValueError

{name}'s input cannot be varying across the axis_name provi

Error message

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

What it means

The unreduced collectives reject inputs that are already marked as 'varying' across the requested axis. Since the collective would reduce across that axis, JAX requires the input not vary along it; this ValueError reports the input aval and the offending axes.

Source

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

  return tree_util.tree_map(
      lambda leaf: unreduced_psum_p.bind(leaf, axes=tuple(axis_name)), x)

unreduced_psum_p = core.Primitive('unreduced_psum')

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

  if any(isinstance(a, int) for a in axes):
    raise ValueError(f'{name} does not accept integer axis_name.'
                     f' Got axis_name={axes}')

  core.check_avals_context_mesh([aval], name)
  check_unreduced_kind(name, aval.mat, out_u_kind)
  out_u = frozenset(u for u in aval.mat.unreduced if u not in axes)
  kind = aval.mat.unreduced_kind if out_u else None
  out_mat = aval.mat.update(unreduced=out_u, unreduced_kind=kind)
  out_aval = aval.update(manual_axis_type=out_mat)
  return out_aval, {core.NamedAxisEffect(axis) for axis in axes}

def _unreduced_psum_abstract_eval(aval, *, axes):
  return _unreduced_psum_pmax_pmin_abstract_eval(
      'unreduced_psum', UnreducedKind.sum, aval, axes=axes)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Remove the axis from the input's varying set (annotate/produce the value as not varying along it) before calling the collective
2. Use a different collective appropriate for varying data (e.g. all_gather plus reduce)
3. Verify `jax.typeof(x).mat.varying` and reconcile it with the axis_name argument

Example fix

// before
jax.lax.unreduced_psum(x, 'data')  # x.varying includes 'data'
// after
jax.lax.unreduced_psum(x, 'data')  # after ensuring x is not varying along 'data', e.g. reshard/annotate accordingly

import jax
assert not (set(jax.typeof(x).mat.varying) & {axis_name}), 'input varies along axis'

def is_invariant_along(x, axis: str) -> bool:
    return axis not in jax.typeof(x).mat.varying

When it happens

Trigger: Calling `unreduced_psum`/`unreduced_pmax`/`unreduced_pmin` with an axis_name that appears in `jax.typeof(x).mat.varying`.

Common situations: Passing a per-device-different (sharded/varying) tensor to a collective expecting a replicated-but-unreduced value; inconsistent sharding annotations after mesh or sharding spec changes.

Related errors

• unbound axis name: {axis_name}
• unreduced_psum_scatter's input cannot be varying across the
• {name} only accepts inputs that are unreduced. Got {aval.str
• {name} is a Unreduced -> Invariant collective. This means th
• preduced is a Invariant->Reduced collective. This means that

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