{"record":{"id":"8cbbec8b4a475441","repo":"jax-ml/jax","slug":"axis-index-groups-not-supported-in-vmap-collective","errorCode":null,"errorMessage":"axis_index_groups not supported in vmap collectives. Please open a feature request!","messagePattern":"axis_index_groups not supported in vmap collectives\\. Please open a feature request!","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/parallel.py","lineNumber":889,"sourceCode":"def _constant_reduction(prim, axis_data, arg, axes, axis_index_groups):\n  assert axis_data.name in axes\n  if axis_index_groups: raise NotImplementedError\n  new_axes = tuple(n for n in axes if n != axis_data.name)\n  if new_axes:\n    arg = (prim.bind(arg, axes=new_axes) if prim is psum_invariant_p else\n           prim.bind(arg, axes=new_axes, axis_index_groups=axis_index_groups))\n  if prim is psum_p:\n    out = lax._const(arg, axis_data.size) * arg\n  elif prim in (pmin_p, pmax_p):\n    out = arg\n  else:\n    raise Exception(f\"Unrecognized reducer: {prim}\")\n  return out, None\n\ndef _reduction_with_positional_batcher(\n    prim, v, d, axis_index_groups, transform_unmapped, transform_mapped):\n  if axis_index_groups is not None:\n    raise NotImplementedError(\"axis_index_groups not supported in vmap collectives. \"\n                              \"Please open a feature request!\")\n  v = v if d is None or d == 0 else _moveaxis(d, 0, v)\n  if d is None:\n    unmapped_axes, unmapped_vals_in = transform_unmapped(0, v)\n    return (prim.bind(unmapped_vals_in, axes=unmapped_axes)\n            if prim is psum_invariant_p else\n            prim.bind(unmapped_vals_in, axes=unmapped_axes, axis_index_groups=None))\n\n  mapped_axes, mapped_vals_in = transform_mapped(0, v)\n  return (prim.bind(mapped_vals_in, axes=mapped_axes)\n          if prim is psum_invariant_p else\n          prim.bind(mapped_vals_in, axes=mapped_axes, axis_index_groups=None))\n\ndef _reduction_batcher(prim, v, d, *, axes, axis_index_groups):\n  assert not prim.multiple_results\n  if not any(isinstance(axis, int) for axis in axes):\n    out = (prim.bind(v, axes=axes) if prim is psum_invariant_p else\n           prim.bind(v, axes=axes, axis_index_groups=axis_index_groups))","sourceCodeStart":871,"sourceCodeEnd":907,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/parallel.py#L871-L907","documentation":"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.","triggerScenarios":"Calling lax.psum, lax.pmax, or lax.pmin with the axis_index_groups argument inside jax.vmap (batched trace).","commonSituations":"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.","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"],"exampleFix":"// before\nout = jax.vmap(lambda x: lax.psum(x, 'i', axis_index_groups=[[0,1],[2,3]]))(x)\n// after\nout = jax.shard_map(lambda x: lax.psum(x, 'i', axis_index_groups=[[0,1],[2,3]]))(x)  # or drop groups","handlingStrategy":"validation","validationCode":"def safe_psum_vmap(f, x, **kw):\n    assert kw.get('axis_index_groups') is None, 'axis_index_groups unsupported under vmap'\n    return jax.vmap(f)(x)","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Never combine axis_index_groups with vmap-batched collectives","Reserve sub-grouped collectives for shard_map code paths"],"tags":["jax","vmap","collectives","not-implemented"],"backgroundTag":"unsupported-operation","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}