jax-ml/jax · error · ValueError

axis_index_groups must all be the same size for TPU lowering

Error message

axis_index_groups must all be the same size for TPU lowering

What it means

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.

Source

Thrown at jax/_src/lax/parallel.py:1073

  with ir.InsertionPoint.at_block_begin(ctx.module_context.module.body):
    reducer = func_dialect.FuncOp(
        f"{prim.name}_{scalar_aval.dtype}_reducer",
        reducer_type,
    )
  reducer.attributes["sym_visibility"] = ir.StringAttr.get("private")
  ctx.module_context.symbol_table.insert(reducer)
  entry_block = reducer.add_entry_block()
  _lower_reducer_into_block(ctx, prim, scalar_aval, entry_block)
  return reducer


def _all_reduce_lowering(prim, pos_fn, ctx, arg, *, axes, axis_index_groups,
                         is_async=False):
  aval_in, = ctx.avals_in
  if axis_index_groups is not None and ("tpu" in ctx.module_context.platforms):
    len_0 = len(axis_index_groups[0])
    if any(len(g) != len_0 for g in axis_index_groups):
      raise ValueError("axis_index_groups must all be the same size for TPU lowering")
  named_axes, positional_axes = axes_partition = [], []
  for axis in axes:
    axes_partition[isinstance(axis, int)].append(axis)

  if positional_axes:
    reducer = mlir.lower_fun(pos_fn, multiple_results=False)
    def _positional_reduce(aval, arg):
      aval_out = aval.update(
          shape=np.delete(np.array(aval.shape, dtype=np.int64),
                          positional_axes))
      reducer_ctx = ctx.replace(primitive=None, avals_in=[aval], avals_out=[aval_out])
      out, = reducer(reducer_ctx, arg, axes=tuple(positional_axes))
      return out
    arg = _positional_reduce(aval_in, arg)
  if not named_axes:
    return [arg]

  replica_groups = _try_mesh_axes_replica_group(

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Pad or regroup axis_index_groups so every group has the same size
2. Use equally-divisible mesh axis sizes
3. Run on GPU/CPU if ragged grouping is semantically required

Example fix

// before
lax.psum(x, 'i', axis_index_groups=[[0,1],[2]])
// after
lax.psum(x, 'i', axis_index_groups=[[0,1],[2,3]])

def check_groups(groups):
    sizes = {len(g) for g in groups}
    assert len(sizes) == 1, f'groups must be equally sized, got sizes {sizes}'

When it happens

Trigger: Calling psum/pmax/pmin with axis_index_groups of unequal lengths while compiling for TPU.

Common situations: Handling a mesh whose axis size doesn't divide evenly into groups; code that worked on GPU (different lowering) failing on TPU.

Related errors

• Unsupported core type: {core_type}
• run_scoped_p with collective axes is not supported
• __dlpack__ device only supported for TPU pinned host memory
• multi-platform lowering for buffer_callback
• `buffer_callback` not supported on {platform} backend.

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/60e700b2cee96c72. Report an issue: GitHub.