{"record":{"id":"748f886502e72d72","repo":"jax-ml/jax","slug":"please-open-a-feature-request","errorCode":null,"errorMessage":"Please open a feature request!","messagePattern":"Please open a feature request!","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/parallel.py","lineNumber":1502,"sourceCode":"def _all_to_all_batcher(vals_in, dims_in, *, axis_name, split_axis, concat_axis, axis_index_groups,\n                        tiled):\n  x, = vals_in\n  d, = dims_in\n  result = all_to_all_p.bind(\n      x,\n      axis_name=axis_name,\n      split_axis=split_axis + (d <= split_axis),\n      concat_axis=concat_axis + (d <= concat_axis),\n      axis_index_groups=axis_index_groups,\n      tiled=tiled,\n  )\n  return result, d\n\ndef _all_to_all_batched_collective(axis_data, vals_in, dims_in,\n                                   axis_name, split_axis, concat_axis,\n                                   axis_index_groups, tiled):\n  if axis_index_groups is not None:\n    raise NotImplementedError(\"Please open a feature request!\")\n  x, = vals_in\n  d, = dims_in\n  axis_size, frame_name = axis_data.size, axis_data.name\n  axes_names = axis_name if isinstance(axis_name, (list, tuple)) else [axis_name]\n  if d is None and frame_name not in axes_names:\n    out = all_to_all_p.bind(\n        x, axis_name=axis_name, split_axis=split_axis, concat_axis=concat_axis,\n        axis_index_groups=axis_index_groups, tiled=tiled)\n    return out, None\n  if frame_name not in axes_names:\n    return _all_to_all_batcher(\n      vals_in, dims_in, axis_name=axis_name, split_axis=split_axis,\n      concat_axis=concat_axis, axis_index_groups=axis_index_groups, tiled=tiled)\n\n  if d is None:\n    # TODO(sharadmv,apaszke): Remove this broadcast that comes from\n    # all_gather_transpose and instead avoid using all_to_all in\n    # all_gather_transpose.","sourceCodeStart":1484,"sourceCodeEnd":1520,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/parallel.py#L1484-L1520","documentation":"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.","triggerScenarios":"jax.vmap over a call to lax.all_to_all(..., axis_index_groups=...).","commonSituations":"Vectorizing sub-group all-to-all communication patterns; porting shard_map pipelines into vmap.","solutions":["Remove axis_index_groups and use full-axis all_to_all under vmap","Use jax.shard_map for sub-grouped all_to_all","File an upstream feature request"],"exampleFix":"// before\njax.vmap(lambda x: lax.all_to_all(x, 'i', 0, 0, axis_index_groups=grp))(x)\n// after\njax.shard_map(lambda x: lax.all_to_all(x, 'i', 0, 0, axis_index_groups=grp), mesh)(x)","handlingStrategy":"validation","validationCode":"assert axis_index_groups is None, 'all_to_all with axis_index_groups cannot be vmap-batched'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Use shard_map for grouped all_to_all"],"tags":["jax","all-to-all","vmap","not-implemented"],"backgroundTag":"unsupported-operation","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}