jax-ml/jax · error · ValueError

{name} does not accept integer axis_name. Got axis_name={axe

Error message

{name} does not accept integer axis_name. Got axis_name={axes}

What it means

Unlike some collectives that accept integer positional axis indices, the unreduced p-sum/max/min collectives only accept named axes (strings). This ValueError fires when any element of the axis_name argument is an int.

Source

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

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)
unreduced_psum_p.def_effectful_abstract_eval(_unreduced_psum_abstract_eval)

def _unreduced_psum_lowering(ctx, arg, *, axes):
  return _all_reduce_lowering(lax.add_p, lax.reduce_sum, ctx, arg,
                             axes=axes, axis_index_groups=None)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Pass the named mesh axis (string) instead of an integer
2. If you meant a positional array axis, use a different API (e.g. `jax.lax.psum` with a name bound via mesh, or plain `jnp.sum(x, axis=0)`)

Example fix

// before
jax.lax.unreduced_psum(x, 0)
// after
jax.lax.unreduced_psum(x, 'data')

assert all(isinstance(a, str) for a in jax.tree.leaves(axis_name)), 'axis_name must be strings, not ints'

def is_named_axes(axis_name) -> bool:
    return all(isinstance(a, str) for a in (axis_name if isinstance(axis_name, (tuple, list)) else (axis_name,)))

When it happens

Trigger: Calling `unreduced_psum(x, 0)` or passing an integer axis (e.g. from `range(x.ndim)`) instead of a named axis string.

Common situations: Copy-pasting code that uses `psum(x, axis=0)`-style positional axes; loops that build axis lists from integers; refactoring from lax.reduce-style APIs that use axis indices.

Related errors

• unbound axis name: {axis_name}
• unbound axis name: {axis_name}
• all_gather_reduced is a Varying -> Reduced collective. This
• unreduced_psum_scatter is a Unreduced -> Varying collective.
• {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/6ae637af903947b2. Report an issue: GitHub.