jax-ml/jax · error · ValueError

axis_index_groups can only be used with reductions over name

Error message

axis_index_groups can only be used with reductions over named axes, but got: {axes}

What it means

axis_index_groups partitions a named mesh axis into sub-groups; it has no meaning for reductions over positional (integer) axes. The abstract eval of psum-style collectives rejects any mix of axis_index_groups with integer axis indices.

Source

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

      )
  )


def _allreduce_impl(prim, pos_reducer, arg, *, axes, axis_index_groups):
  assert axis_index_groups is None
  if not all(isinstance(axis, int) for axis in axes):
     return dispatch.apply_primitive(prim, arg, axes=axes,
                                     axis_index_groups=axis_index_groups)
  assert all(isinstance(axis, int) for axis in axes)
  return pos_reducer(arg, axes)

def _allreduce_effectful_abstract_eval(aval, *, axes, axis_index_groups):
  _check_axis_names(axes, 'psum')
  named_axes = tuple(axis for axis in axes if not isinstance(axis, int))
  pos_axes = tuple(axis for axis in axes if isinstance(axis, int))
  if axis_index_groups is not None:
    if len(pos_axes) != 0:
      raise ValueError(f"axis_index_groups can only be used with reductions over "
                       f"named axes, but got: {axes}")
  core.check_avals_context_mesh([aval], 'psum')
  check_unreduced_args([aval], axes, 'psum')
  out_aval = ShapedArray(
      lax._reduce_op_shape_rule(aval, axes=pos_axes), aval.dtype,
      sharding=lax._reduce_op_sharding_rule(aval, axes=pos_axes))
  return out_aval, {core.NamedAxisEffect(axis) for axis in named_axes}

# TODO(yashkatariya): Replace this with _psum_invariant_abstract_eval
def _pmin_pmax_abstract_eval(name, aval, *, axes, axis_index_groups):
  if not config._check_vma.value:
    return _allreduce_effectful_abstract_eval(
        aval, axes=axes, axis_index_groups=axis_index_groups)
  return _psum_invariant_abstract_eval(name, aval, axes=axes)

def _check_axis_names(axes, api_name):
  named_axes = tuple(axis for axis in axes if not isinstance(axis, int))
  axis_env = core.get_axis_env()

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Pass a named axis string (declared in shard_map/vmap axis_name) instead of an int
2. Drop axis_index_groups if you don't need sub-groups
3. Reduce positional axes with lax.sum/etc. instead of collectives

Example fix

// before
lax.psum(x, 1, axis_index_groups=[[0,1],[2,3]])
// after
jax.shard_map(lambda x: lax.psum(x, 'i', axis_index_groups=[[0,1],[2,3]]), mesh)(x)

def check_psum_args(axes, groups):
    if groups is not None and any(isinstance(a, int) for a in (axes if isinstance(axes, (list,tuple)) else [axes])):
        raise ValueError('axis_index_groups requires named axes only')

When it happens

Trigger: Calling lax.psum/pmax/pmin (or psum0) with axis_index_groups and passing an int axis (e.g. psum(x, 0, axis_index_groups=...)).

Common situations: Converting code that reduced positional axes to named-axis collectives while keeping the int axis; copy-pasting shard_map examples with wrong axis argument type.

Related errors

• unbound axis name: {axis_name}
• {name} cannot accept args which are unreduced. Got {a.str_sh
• {name} cannot accept args with unreduced_kind={a.mat.unreduc
• {name} cannot accept args which are reduced. Got {a.str_shor
• unbound axis name: {axis_name}

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