jax-ml/jax · error · NotImplementedError

axis_index_groups not supported in vmap collectives. Please

Error message

axis_index_groups not supported in vmap collectives. Please open a feature request!

What it means

Raised when a collective reduction (psum/pmax/pmin) with axis_index_groups is used inside vmap. The vmap batching rule for reduction collectives does not implement support for axis_index_groups, so it raises NotImplementedError instead of silently producing wrong results.

Source

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

def _constant_reduction(prim, axis_data, arg, axes, axis_index_groups):
  assert axis_data.name in axes
  if axis_index_groups: raise NotImplementedError
  new_axes = tuple(n for n in axes if n != axis_data.name)
  if new_axes:
    arg = (prim.bind(arg, axes=new_axes) if prim is psum_invariant_p else
           prim.bind(arg, axes=new_axes, axis_index_groups=axis_index_groups))
  if prim is psum_p:
    out = lax._const(arg, axis_data.size) * arg
  elif prim in (pmin_p, pmax_p):
    out = arg
  else:
    raise Exception(f"Unrecognized reducer: {prim}")
  return out, None

def _reduction_with_positional_batcher(
    prim, v, d, axis_index_groups, transform_unmapped, transform_mapped):
  if axis_index_groups is not None:
    raise NotImplementedError("axis_index_groups not supported in vmap collectives. "
                              "Please open a feature request!")
  v = v if d is None or d == 0 else _moveaxis(d, 0, v)
  if d is None:
    unmapped_axes, unmapped_vals_in = transform_unmapped(0, v)
    return (prim.bind(unmapped_vals_in, axes=unmapped_axes)
            if prim is psum_invariant_p else
            prim.bind(unmapped_vals_in, axes=unmapped_axes, axis_index_groups=None))

  mapped_axes, mapped_vals_in = transform_mapped(0, v)
  return (prim.bind(mapped_vals_in, axes=mapped_axes)
          if prim is psum_invariant_p else
          prim.bind(mapped_vals_in, axes=mapped_axes, axis_index_groups=None))

def _reduction_batcher(prim, v, d, *, axes, axis_index_groups):
  assert not prim.multiple_results
  if not any(isinstance(axis, int) for axis in axes):
    out = (prim.bind(v, axes=axes) if prim is psum_invariant_p else
           prim.bind(v, axes=axes, axis_index_groups=axis_index_groups))

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Remove axis_index_groups and rely on named-axis reduction under vmap
Use jax.shard_map instead of vmap for sub-group collective patterns
File a feature request at github.com/google/jax/issues as the message suggests

Example fix

// before
out = jax.vmap(lambda x: lax.psum(x, 'i', axis_index_groups=[[0,1],[2,3]]))(x)
// after
out = jax.shard_map(lambda x: lax.psum(x, 'i', axis_index_groups=[[0,1],[2,3]]))(x)  # or drop groups

Defensive patterns

Strategy: validation

Validate before calling

def safe_psum_vmap(f, x, **kw):
    assert kw.get('axis_index_groups') is None, 'axis_index_groups unsupported under vmap'
    return jax.vmap(f)(x)

Prevention

Never combine axis_index_groups with vmap-batched collectives
Reserve sub-grouped collectives for shard_map code paths

When it happens

Trigger: Calling lax.psum, lax.pmax, or lax.pmin with the axis_index_groups argument inside jax.vmap (batched trace).

Common situations: Emulating sub-group collective ops (e.g. grouped sums within a partition of devices) while also vectorizing with vmap; migrating shard_map code into vmap.

Understand the failure class

Background: UnsupportedOperationException and "is not supported" errors: when a library deliberately refuses a call — this error's family across 30 libraries.

Related errors

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