{"record":{"id":"8526efe030c38f62","repo":"jax-ml/jax","slug":"axis-index-groups-only-supported-for-sums-over-jus","errorCode":null,"errorMessage":"axis_index_groups only supported for sums over just named axes","messagePattern":"axis_index_groups only supported for sums over just named axes","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/parallel.py","lineNumber":161,"sourceCode":"    return x\n  def bind(leaf):\n    from_ = _get_from(core.typeof(leaf), axes, 'jax.lax.psum')\n    if from_ == 'unreduced':\n      if axis_index_groups is not None:\n        raise NotImplementedError\n      return unreduced_psum(leaf, axes)\n    else:\n      return _psum(leaf, axes, axis_index_groups=axis_index_groups,\n                   is_async=is_async)\n  return tree_util.tree_map(bind, x)\n\ndef _psum(x, axis_name, *, axis_index_groups, is_async):\n  if not isinstance(axis_name, (tuple, list)):\n    axis_name = (axis_name,)\n  if not axis_name:\n    return x\n  if any(isinstance(axis, int) for axis in axis_name) and axis_index_groups is not None:\n    raise ValueError(\"axis_index_groups only supported for sums over just named axes\")\n  _validate_reduce_axis_index_groups(axis_index_groups)\n  leaves, treedef = tree_util.tree_flatten(x)\n  leaves = [lax.convert_element_type(l, np.int32)\n            if dtypes.dtype(l) == np.bool_ else l for l in leaves]\n  axis_index_groups = _canonicalize_axis_index_groups(axis_index_groups)\n  # handle the constant case specially\n  if all(not isinstance(leaf, core.Tracer) for leaf in leaves):\n    named_axes, pos_axes = axes_partition = [], []\n    for axis in axis_name:\n      axes_partition[isinstance(axis, int)].append(axis)\n    def pos_reduce(x):\n      if not pos_axes:\n        return x\n      return lax.reduce_sum(x, [canonicalize_axis(axis, getattr(x, 'ndim', 0))\n                                for axis in pos_axes])\n    if axis_index_groups is not None:\n      assert not pos_axes\n      size = len(axis_index_groups[0])","sourceCodeStart":143,"sourceCodeEnd":179,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/parallel.py#L143-L179","documentation":"jax.lax.psum's axis_name may contain named axes (strings/hashables) or positional integer axes, but axis_index_groups (grouping device indices into subgroups) is only defined when reducing purely over named axes. Mixing an integer positional axis with axis_index_groups raises this ValueError.","triggerScenarios":"psum(x, (0,), axis_index_groups=[[0,1]]) or psum(x, 1, axis_index_groups=...) — any call where the axis tuple contains an int and axis_index_groups is not None.","commonSituations":"Converting vmap axes to pmap-style collectives and reusing integer axis indices; adding grouped reductions to code that reduces over a batch dimension by position.","solutions":["Replace integer axes with the corresponding named axis from your pmap/shard_map declaration","If you intended a positional reduction, use jnp.sum(x, axis=int) instead of psum","Register a name for the axis (pmap(axis_name='i') / spmd axes) and pass that name"],"exampleFix":"// before\ny = jax.lax.psum(x, 0, axis_index_groups=[[0,1],[2,3]])\n\n// after\ny = jax.lax.psum(x, 'i', axis_index_groups=[[0,1],[2,3]])  # pmap axis_name='i'","handlingStrategy":"validation","validationCode":"axes = axis_name if isinstance(axis_name, (tuple, list)) else (axis_name,)\nassert not (axis_index_groups is not None and any(isinstance(a, int) for a in axes))","typeGuard":"def named_axes_only(axis_name):\n    axes = axis_name if isinstance(axis_name, (tuple, list)) else (axis_name,)\n    assert all(not isinstance(a, int) for a in axes), 'psum needs named axes for groups'","tryCatchPattern":null,"preventionTips":["Always use named axes for collectives","Lint against integer axis args in collective wrappers"],"tags":["jax","psum","axis","validation"],"backgroundTag":"invalid-axis-argument","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}