jax-ml/jax · error · RuntimeError

all_to_all must be used within a mapped context like vmap or

Error message

all_to_all must be used within a mapped context like vmap or shard_map.

What it means

all_to_all has no eager top-level implementation: it only has meaning inside a mapped context (vmap with axis_name or shard_map) that binds the axis name. Calling it directly (or via a plain jit without mapping) hits the default impl which raises RuntimeError.

Source

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

  check_unreduced_args([input_aval], axis_name, 'all_to_all')
  shape = list(input_aval.shape)
  axis_size = (
      _axis_size(axis_name)
      if axis_index_groups is None
      else len(axis_index_groups[0])
  )
  assert shape[split_axis] % axis_size == 0, (shape[split_axis], axis_size)
  shape[split_axis] //= axis_size
  shape[concat_axis] *= axis_size
  vma = collective_vma_rule('all_to_all', axis_name, input_aval)
  out_aval = input_aval.update(
    shape=tuple(shape), weak_type=False,
    manual_axis_type=input_aval.mat.update(varying=vma))
  effects = {*map(core.NamedAxisEffect, axis_name)}
  return out_aval, effects

def _all_to_all_impl(*args, **kwargs):
  raise RuntimeError("all_to_all must be used within a mapped context"
                     " like vmap or shard_map.")

all_to_all_p = core.Primitive('all_to_all')
all_to_all_p.def_impl(_all_to_all_impl)
all_to_all_p.def_effectful_abstract_eval(_all_to_all_effectful_abstract_eval)
mlir.register_lowering(all_to_all_p, _all_to_all_lowering)
ad.deflinear2(all_to_all_p, _all_to_all_transpose_rule)
batching.fancy_primitive_batchers[all_to_all_p] = _all_to_all_batched_collective


def _ragged_all_to_all_lowering(
    ctx, operand, output, input_offsets, send_sizes, output_offsets, recv_sizes,
    *, axis_name, axis_index_groups
):
  replica_groups = _replica_groups(ctx.module_context.axis_context, axis_name,
                                   axis_index_groups)

  # Assumes all groups are the same size

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Wrap in jax.vmap(f, axis_name='i') or jax.shard_map(f, mesh)
2. Ensure the mapped axis matches the all_to_all axis_name
3. Use non-collective ops (reshape/transpose) for single-device equivalents

Example fix

// before
y = lax.all_to_all(x, 'i', 0, 0)
// after
y = jax.vmap(lambda b: lax.all_to_all(b, 'i', 0, 0), axis_name='i')(x)

null  # structural: ensure all_to_all appears only inside vmap(axis_name=...) or shard_map

try:
    y = lax.all_to_all(x, 'i', 0, 0)
except RuntimeError as e:
    if 'mapped context' in str(e):
        y = jax.vmap(lambda b: lax.all_to_all(b, 'i', 0, 0), axis_name='i')(x)
    else: raise

When it happens

Trigger: Calling jax.lax.all_to_all outside vmap/shard_map, e.g. eagerly or under jax.jit without shard_map.

Common situations: Quick REPL experiments with collectives; refactoring that accidentally removed the vmap/shard_map wrapper.

Related errors

• The size of all_to_all split_axis ({x.shape[split_axis]}) ha
• all_to_all requires the size of the mapped axis axis_name to
• Replica groups must be equally sized
• Please open a feature request!
• numpy masked arrays are not supported as direct inputs to JA

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