{"record":{"id":"32393d86f098a9ab","repo":"jax-ml/jax","slug":"axis-index-groups-must-cover-all-indices-exactly-o","errorCode":null,"errorMessage":"axis_index_groups must cover all indices exactly once","messagePattern":"axis_index_groups must cover all indices exactly once","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/parallel.py","lineNumber":341,"sourceCode":"def pargmax(x, axis_name):\n  if isinstance(axis_name, (tuple, list)):\n    raise TypeError(f\"pargmin only accepts a single axis, got {axis_name}\")\n  return _axis_index_of_val(x, pmax(x, axis_name), axis_name)\n\ndef _axis_index_of_val(x, val, axis_name):\n  idx = axis_index(axis_name)\n  mask = (val == x)\n  validx = lax.select(mask,\n                      lax.full(mask.shape, idx),\n                      lax.full(mask.shape, dtypes.iinfo(idx.dtype).max, idx.dtype))\n  return pmin(validx, axis_name)\n\ndef _validate_reduce_axis_index_groups(axis_index_groups):\n  if axis_index_groups is None:\n    return\n  axis_space = range(sum(len(group) for group in axis_index_groups))\n  if {i for g in axis_index_groups for i in g} != set(axis_space):\n    raise ValueError(\"axis_index_groups must cover all indices exactly once\")\n\ndef _canonicalize_axis_index_groups(axis_index_groups):\n  if axis_index_groups is None:\n    return\n  return tuple(map(tuple, axis_index_groups))\n\n\ndef pbroadcast(x, axis_name, source):\n  \"\"\"Perform a collective broadcast and replicate from ``source``.\n\n  This is equivalent to::\n\n    def pbroadcast(x, axis_name, source):\n      masked = jnp.where(axis_index(axis_name) == source, x, zeros_like(x))\n      return psum(masked, axis_name)\n\n  but implemented in a hardware optimized way.\n","sourceCodeStart":323,"sourceCodeEnd":359,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/parallel.py#L323-L359","documentation":"When using axis_index_groups with psum/pmax/pmin, the groups must partition all device indices of the mapped axis: every index from 0..total-1 must appear in exactly one group. Duplicate or missing indices fail this set-equality check.","triggerScenarios":"axis_index_groups=[[0,1],[2,2]] (duplicate 2, missing 3) or [[0,1]] on a 4-device axis (missing 2,3).","commonSituations":"Hand-written subgroups that forget the last device; off-by-one when generating groups programmatically; changing device count without regenerating groups.","solutions":["Generate groups programmatically covering range(axis_size), e.g. reshape of arange","Validate coverage before calling: set(flat) == set(range(total)) and len(flat) == total","Recompute groups whenever the mesh/axis size changes"],"exampleFix":"# before\naxis_index_groups=[[0,1],[2,2]]\ny = jax.lax.psum(x, 'i', axis_index_groups=axis_index_groups)\n\n# after\naxis_index_groups=[[0,1],[2,3]]\ny = jax.lax.psum(x, 'i', axis_index_groups=axis_index_groups)","handlingStrategy":"validation","validationCode":"flat = [i for g in axis_index_groups for i in g]\nassert sorted(flat) == list(range(len(flat))), 'groups must cover indices exactly once'","typeGuard":"def valid_groups(groups, axis_size):\n    flat = [i for g in groups for i in g]\n    return sorted(flat) == list(range(axis_size))","tryCatchPattern":null,"preventionTips":["Generate groups from arange(axis_size).reshape(-1, g)","Re-validate groups after mesh size changes"],"tags":["jax","psum","validation","axis-index-groups"],"backgroundTag":"invalid-group-partition","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}