{"record":{"id":"3c4d033d00b33c38","repo":"jax-ml/jax","slug":"axis-index-groups-can-only-be-used-with-reductions","errorCode":null,"errorMessage":"axis_index_groups can only be used with reductions over named axes, but got: {axes}","messagePattern":"axis_index_groups can only be used with reductions over named axes, but got: (.+?)","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/parallel.py","lineNumber":1012,"sourceCode":"      )\n  )\n\n\ndef _allreduce_impl(prim, pos_reducer, arg, *, axes, axis_index_groups):\n  assert axis_index_groups is None\n  if not all(isinstance(axis, int) for axis in axes):\n     return dispatch.apply_primitive(prim, arg, axes=axes,\n                                     axis_index_groups=axis_index_groups)\n  assert all(isinstance(axis, int) for axis in axes)\n  return pos_reducer(arg, axes)\n\ndef _allreduce_effectful_abstract_eval(aval, *, axes, axis_index_groups):\n  _check_axis_names(axes, 'psum')\n  named_axes = tuple(axis for axis in axes if not isinstance(axis, int))\n  pos_axes = tuple(axis for axis in axes if isinstance(axis, int))\n  if axis_index_groups is not None:\n    if len(pos_axes) != 0:\n      raise ValueError(f\"axis_index_groups can only be used with reductions over \"\n                       f\"named axes, but got: {axes}\")\n  core.check_avals_context_mesh([aval], 'psum')\n  check_unreduced_args([aval], axes, 'psum')\n  out_aval = ShapedArray(\n      lax._reduce_op_shape_rule(aval, axes=pos_axes), aval.dtype,\n      sharding=lax._reduce_op_sharding_rule(aval, axes=pos_axes))\n  return out_aval, {core.NamedAxisEffect(axis) for axis in named_axes}\n\n# TODO(yashkatariya): Replace this with _psum_invariant_abstract_eval\ndef _pmin_pmax_abstract_eval(name, aval, *, axes, axis_index_groups):\n  if not config._check_vma.value:\n    return _allreduce_effectful_abstract_eval(\n        aval, axes=axes, axis_index_groups=axis_index_groups)\n  return _psum_invariant_abstract_eval(name, aval, axes=axes)\n\ndef _check_axis_names(axes, api_name):\n  named_axes = tuple(axis for axis in axes if not isinstance(axis, int))\n  axis_env = core.get_axis_env()","sourceCodeStart":994,"sourceCodeEnd":1030,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/parallel.py#L994-L1030","documentation":"axis_index_groups partitions a named mesh axis into sub-groups; it has no meaning for reductions over positional (integer) axes. The abstract eval of psum-style collectives rejects any mix of axis_index_groups with integer axis indices.","triggerScenarios":"Calling lax.psum/pmax/pmin (or psum0) with axis_index_groups and passing an int axis (e.g. psum(x, 0, axis_index_groups=...)).","commonSituations":"Converting code that reduced positional axes to named-axis collectives while keeping the int axis; copy-pasting shard_map examples with wrong axis argument type.","solutions":["Pass a named axis string (declared in shard_map/vmap axis_name) instead of an int","Drop axis_index_groups if you don't need sub-groups","Reduce positional axes with lax.sum/etc. instead of collectives"],"exampleFix":"// before\nlax.psum(x, 1, axis_index_groups=[[0,1],[2,3]])\n// after\njax.shard_map(lambda x: lax.psum(x, 'i', axis_index_groups=[[0,1],[2,3]]), mesh)(x)","handlingStrategy":"validation","validationCode":"def check_psum_args(axes, groups):\n    if groups is not None and any(isinstance(a, int) for a in (axes if isinstance(axes, (list,tuple)) else [axes])):\n        raise ValueError('axis_index_groups requires named axes only')","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Use string axis names whenever passing axis_index_groups"],"tags":["jax","collectives","axis-validation"],"backgroundTag":"invalid-argument","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}