{"record":{"id":"3bfcc10ee360159c","repo":"jax-ml/jax","slug":"all-to-all-must-be-used-within-a-mapped-context-li","errorCode":null,"errorMessage":"all_to_all must be used within a mapped context like vmap or shard_map.","messagePattern":"all_to_all must be used within a mapped context like vmap or shard_map\\.","errorType":"exception","errorClass":"RuntimeError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/parallel.py","lineNumber":1597,"sourceCode":"  check_unreduced_args([input_aval], axis_name, 'all_to_all')\n  shape = list(input_aval.shape)\n  axis_size = (\n      _axis_size(axis_name)\n      if axis_index_groups is None\n      else len(axis_index_groups[0])\n  )\n  assert shape[split_axis] % axis_size == 0, (shape[split_axis], axis_size)\n  shape[split_axis] //= axis_size\n  shape[concat_axis] *= axis_size\n  vma = collective_vma_rule('all_to_all', axis_name, input_aval)\n  out_aval = input_aval.update(\n    shape=tuple(shape), weak_type=False,\n    manual_axis_type=input_aval.mat.update(varying=vma))\n  effects = {*map(core.NamedAxisEffect, axis_name)}\n  return out_aval, effects\n\ndef _all_to_all_impl(*args, **kwargs):\n  raise RuntimeError(\"all_to_all must be used within a mapped context\"\n                     \" like vmap or shard_map.\")\n\nall_to_all_p = core.Primitive('all_to_all')\nall_to_all_p.def_impl(_all_to_all_impl)\nall_to_all_p.def_effectful_abstract_eval(_all_to_all_effectful_abstract_eval)\nmlir.register_lowering(all_to_all_p, _all_to_all_lowering)\nad.deflinear2(all_to_all_p, _all_to_all_transpose_rule)\nbatching.fancy_primitive_batchers[all_to_all_p] = _all_to_all_batched_collective\n\n\ndef _ragged_all_to_all_lowering(\n    ctx, operand, output, input_offsets, send_sizes, output_offsets, recv_sizes,\n    *, axis_name, axis_index_groups\n):\n  replica_groups = _replica_groups(ctx.module_context.axis_context, axis_name,\n                                   axis_index_groups)\n\n  # Assumes all groups are the same size","sourceCodeStart":1579,"sourceCodeEnd":1615,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/parallel.py#L1579-L1615","documentation":"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.","triggerScenarios":"Calling jax.lax.all_to_all outside vmap/shard_map, e.g. eagerly or under jax.jit without shard_map.","commonSituations":"Quick REPL experiments with collectives; refactoring that accidentally removed the vmap/shard_map wrapper.","solutions":["Wrap in jax.vmap(f, axis_name='i') or jax.shard_map(f, mesh)","Ensure the mapped axis matches the all_to_all axis_name","Use non-collective ops (reshape/transpose) for single-device equivalents"],"exampleFix":"// before\ny = lax.all_to_all(x, 'i', 0, 0)\n// after\ny = jax.vmap(lambda b: lax.all_to_all(b, 'i', 0, 0), axis_name='i')(x)","handlingStrategy":"validation","validationCode":"null  # structural: ensure all_to_all appears only inside vmap(axis_name=...) or shard_map","typeGuard":null,"tryCatchPattern":"try:\n    y = lax.all_to_all(x, 'i', 0, 0)\nexcept RuntimeError as e:\n    if 'mapped context' in str(e):\n        y = jax.vmap(lambda b: lax.all_to_all(b, 'i', 0, 0), axis_name='i')(x)\n    else: raise","preventionTips":["Always test collectives inside their mapping wrapper"],"tags":["jax","all-to-all","mapped-context"],"backgroundTag":"missing-context","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}