jax-ml/jax · error · ValueError

axis_index_groups must cover all indices exactly once

Error message

axis_index_groups must cover all indices exactly once

What it means

When using axis_index_groups with psum/pmax/pmin, the groups must partition all device indices of the mapped axis: every index from 0..total-1 must appear in exactly one group. Duplicate or missing indices fail this set-equality check.

Source

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

def pargmax(x, axis_name):
  if isinstance(axis_name, (tuple, list)):
    raise TypeError(f"pargmin only accepts a single axis, got {axis_name}")
  return _axis_index_of_val(x, pmax(x, axis_name), axis_name)

def _axis_index_of_val(x, val, axis_name):
  idx = axis_index(axis_name)
  mask = (val == x)
  validx = lax.select(mask,
                      lax.full(mask.shape, idx),
                      lax.full(mask.shape, dtypes.iinfo(idx.dtype).max, idx.dtype))
  return pmin(validx, axis_name)

def _validate_reduce_axis_index_groups(axis_index_groups):
  if axis_index_groups is None:
    return
  axis_space = range(sum(len(group) for group in axis_index_groups))
  if {i for g in axis_index_groups for i in g} != set(axis_space):
    raise ValueError("axis_index_groups must cover all indices exactly once")

def _canonicalize_axis_index_groups(axis_index_groups):
  if axis_index_groups is None:
    return
  return tuple(map(tuple, axis_index_groups))


def pbroadcast(x, axis_name, source):
  """Perform a collective broadcast and replicate from ``source``.

  This is equivalent to::

    def pbroadcast(x, axis_name, source):
      masked = jnp.where(axis_index(axis_name) == source, x, zeros_like(x))
      return psum(masked, axis_name)

  but implemented in a hardware optimized way.

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Generate groups programmatically covering range(axis_size), e.g. reshape of arange
Validate coverage before calling: set(flat) == set(range(total)) and len(flat) == total
Recompute groups whenever the mesh/axis size changes

Example fix

# before
axis_index_groups=[[0,1],[2,2]]
y = jax.lax.psum(x, 'i', axis_index_groups=axis_index_groups)

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

Defensive patterns

Strategy: validation

Validate before calling

flat = [i for g in axis_index_groups for i in g]
assert sorted(flat) == list(range(len(flat))), 'groups must cover indices exactly once'

Type guard

def valid_groups(groups, axis_size):
    flat = [i for g in groups for i in g]
    return sorted(flat) == list(range(axis_size))

Prevention

Generate groups from arange(axis_size).reshape(-1, g)
Re-validate groups after mesh size changes

When it happens

Trigger: axis_index_groups=[[0,1],[2,2]] (duplicate 2, missing 3) or [[0,1]] on a 4-device axis (missing 2,3).

Common situations: Hand-written subgroups that forget the last device; off-by-one when generating groups programmatically; changing device count without regenerating groups.

Related errors

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