{"record":{"id":"60e700b2cee96c72","repo":"jax-ml/jax","slug":"axis-index-groups-must-all-be-the-same-size-for-tp","errorCode":null,"errorMessage":"axis_index_groups must all be the same size for TPU lowering","messagePattern":"axis_index_groups must all be the same size for TPU lowering","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/parallel.py","lineNumber":1073,"sourceCode":"  with ir.InsertionPoint.at_block_begin(ctx.module_context.module.body):\n    reducer = func_dialect.FuncOp(\n        f\"{prim.name}_{scalar_aval.dtype}_reducer\",\n        reducer_type,\n    )\n  reducer.attributes[\"sym_visibility\"] = ir.StringAttr.get(\"private\")\n  ctx.module_context.symbol_table.insert(reducer)\n  entry_block = reducer.add_entry_block()\n  _lower_reducer_into_block(ctx, prim, scalar_aval, entry_block)\n  return reducer\n\n\ndef _all_reduce_lowering(prim, pos_fn, ctx, arg, *, axes, axis_index_groups,\n                         is_async=False):\n  aval_in, = ctx.avals_in\n  if axis_index_groups is not None and (\"tpu\" in ctx.module_context.platforms):\n    len_0 = len(axis_index_groups[0])\n    if any(len(g) != len_0 for g in axis_index_groups):\n      raise ValueError(\"axis_index_groups must all be the same size for TPU lowering\")\n  named_axes, positional_axes = axes_partition = [], []\n  for axis in axes:\n    axes_partition[isinstance(axis, int)].append(axis)\n\n  if positional_axes:\n    reducer = mlir.lower_fun(pos_fn, multiple_results=False)\n    def _positional_reduce(aval, arg):\n      aval_out = aval.update(\n          shape=np.delete(np.array(aval.shape, dtype=np.int64),\n                          positional_axes))\n      reducer_ctx = ctx.replace(primitive=None, avals_in=[aval], avals_out=[aval_out])\n      out, = reducer(reducer_ctx, arg, axes=tuple(positional_axes))\n      return out\n    arg = _positional_reduce(aval_in, arg)\n  if not named_axes:\n    return [arg]\n\n  replica_groups = _try_mesh_axes_replica_group(","sourceCodeStart":1055,"sourceCodeEnd":1091,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/parallel.py#L1055-L1091","documentation":"On TPU, psum-family collectives with axis_index_groups lower to an AllReduce with explicit replica groups, which XLA requires to be uniformly sized. The lowering checks all groups share the same length before emitting HLO.","triggerScenarios":"Calling psum/pmax/pmin with axis_index_groups of unequal lengths while compiling for TPU.","commonSituations":"Handling a mesh whose axis size doesn't divide evenly into groups; code that worked on GPU (different lowering) failing on TPU.","solutions":["Pad or regroup axis_index_groups so every group has the same size","Use equally-divisible mesh axis sizes","Run on GPU/CPU if ragged grouping is semantically required"],"exampleFix":"// before\nlax.psum(x, 'i', axis_index_groups=[[0,1],[2]])\n// after\nlax.psum(x, 'i', axis_index_groups=[[0,1],[2,3]])","handlingStrategy":"validation","validationCode":"def check_groups(groups):\n    sizes = {len(g) for g in groups}\n    assert len(sizes) == 1, f'groups must be equally sized, got sizes {sizes}'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Unit-test group size uniformity on TPU-targeted configs"],"tags":["jax","tpu","collectives","lowering"],"backgroundTag":"unequal-group-sizes","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}