jax-ml/jax · error · NotImplementedError

Please open a feature request!

Error message

Please open a feature request!

What it means

The vmap batching rule for all_to_all does not support axis_index_groups, so batching an all_to_all that partitions the axis into sub-groups raises NotImplementedError requesting a feature request.

Source

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

def _all_to_all_batcher(vals_in, dims_in, *, axis_name, split_axis, concat_axis, axis_index_groups,
                        tiled):
  x, = vals_in
  d, = dims_in
  result = all_to_all_p.bind(
      x,
      axis_name=axis_name,
      split_axis=split_axis + (d <= split_axis),
      concat_axis=concat_axis + (d <= concat_axis),
      axis_index_groups=axis_index_groups,
      tiled=tiled,
  )
  return result, d

def _all_to_all_batched_collective(axis_data, vals_in, dims_in,
                                   axis_name, split_axis, concat_axis,
                                   axis_index_groups, tiled):
  if axis_index_groups is not None:
    raise NotImplementedError("Please open a feature request!")
  x, = vals_in
  d, = dims_in
  axis_size, frame_name = axis_data.size, axis_data.name
  axes_names = axis_name if isinstance(axis_name, (list, tuple)) else [axis_name]
  if d is None and frame_name not in axes_names:
    out = all_to_all_p.bind(
        x, axis_name=axis_name, split_axis=split_axis, concat_axis=concat_axis,
        axis_index_groups=axis_index_groups, tiled=tiled)
    return out, None
  if frame_name not in axes_names:
    return _all_to_all_batcher(
      vals_in, dims_in, axis_name=axis_name, split_axis=split_axis,
      concat_axis=concat_axis, axis_index_groups=axis_index_groups, tiled=tiled)

  if d is None:
    # TODO(sharadmv,apaszke): Remove this broadcast that comes from
    # all_gather_transpose and instead avoid using all_to_all in
    # all_gather_transpose.

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Remove axis_index_groups and use full-axis all_to_all under vmap
2. Use jax.shard_map for sub-grouped all_to_all
3. File an upstream feature request

Example fix

// before
jax.vmap(lambda x: lax.all_to_all(x, 'i', 0, 0, axis_index_groups=grp))(x)
// after
jax.shard_map(lambda x: lax.all_to_all(x, 'i', 0, 0, axis_index_groups=grp), mesh)(x)

assert axis_index_groups is None, 'all_to_all with axis_index_groups cannot be vmap-batched'

When it happens

Trigger: jax.vmap over a call to lax.all_to_all(..., axis_index_groups=...).

Common situations: Vectorizing sub-group all-to-all communication patterns; porting shard_map pipelines 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

• ragged_dot vmap over any dim but 0 - NYI
• axis_index_groups not supported in vmap collectives. Please
• pbroadcast batcher only supports a single axis
• Please file an issue at https://github.com/jax-ml/jax/issues
• reduce_window batching is not implemented for initial values

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