jax-ml/jax · error · ValueError

axis_index_groups only supported for sums over just named ax

Error message

axis_index_groups only supported for sums over just named axes

What it means

jax.lax.psum's axis_name may contain named axes (strings/hashables) or positional integer axes, but axis_index_groups (grouping device indices into subgroups) is only defined when reducing purely over named axes. Mixing an integer positional axis with axis_index_groups raises this ValueError.

Source

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

    return x
  def bind(leaf):
    from_ = _get_from(core.typeof(leaf), axes, 'jax.lax.psum')
    if from_ == 'unreduced':
      if axis_index_groups is not None:
        raise NotImplementedError
      return unreduced_psum(leaf, axes)
    else:
      return _psum(leaf, axes, axis_index_groups=axis_index_groups,
                   is_async=is_async)
  return tree_util.tree_map(bind, x)

def _psum(x, axis_name, *, axis_index_groups, is_async):
  if not isinstance(axis_name, (tuple, list)):
    axis_name = (axis_name,)
  if not axis_name:
    return x
  if any(isinstance(axis, int) for axis in axis_name) and axis_index_groups is not None:
    raise ValueError("axis_index_groups only supported for sums over just named axes")
  _validate_reduce_axis_index_groups(axis_index_groups)
  leaves, treedef = tree_util.tree_flatten(x)
  leaves = [lax.convert_element_type(l, np.int32)
            if dtypes.dtype(l) == np.bool_ else l for l in leaves]
  axis_index_groups = _canonicalize_axis_index_groups(axis_index_groups)
  # handle the constant case specially
  if all(not isinstance(leaf, core.Tracer) for leaf in leaves):
    named_axes, pos_axes = axes_partition = [], []
    for axis in axis_name:
      axes_partition[isinstance(axis, int)].append(axis)
    def pos_reduce(x):
      if not pos_axes:
        return x
      return lax.reduce_sum(x, [canonicalize_axis(axis, getattr(x, 'ndim', 0))
                                for axis in pos_axes])
    if axis_index_groups is not None:
      assert not pos_axes
      size = len(axis_index_groups[0])

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Replace integer axes with the corresponding named axis from your pmap/shard_map declaration
2. If you intended a positional reduction, use jnp.sum(x, axis=int) instead of psum
3. Register a name for the axis (pmap(axis_name='i') / spmd axes) and pass that name

Example fix

// before
y = jax.lax.psum(x, 0, axis_index_groups=[[0,1],[2,3]])

// after
y = jax.lax.psum(x, 'i', axis_index_groups=[[0,1],[2,3]])  # pmap axis_name='i'

axes = axis_name if isinstance(axis_name, (tuple, list)) else (axis_name,)
assert not (axis_index_groups is not None and any(isinstance(a, int) for a in axes))

def named_axes_only(axis_name):
    axes = axis_name if isinstance(axis_name, (tuple, list)) else (axis_name,)
    assert all(not isinstance(a, int) for a in axes), 'psum needs named axes for groups'

When it happens

Trigger: psum(x, (0,), axis_index_groups=[[0,1]]) or psum(x, 1, axis_index_groups=...) — any call where the axis tuple contains an int and axis_index_groups is not None.

Common situations: Converting vmap axes to pmap-style collectives and reusing integer axis indices; adding grouped reductions to code that reduces over a batch dimension by position.

Related errors

• Invalid axis {axis} for operand shape {operand.shape}
• axis_index_groups must cover all indices exactly once
• {full_name} must be a pytree prefix with bool leaves or a tu
• Axes mentioned in `manual_axis_type` field of ShapedArray sh
• varying and unreduced cannot have common mesh axes. Got vary

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