jax-ml/jax · error · ValueError

Warp reduction group size should be a power of 2 (got {group

Error message

Warp reduction group size should be a power of 2 (got {group_size})

What it means

Warp reduction (warp_reduce) uses log2(group_size) shuffle iterations, which only works when group_size is a power of two. The preceding assert also requires group_size to divide 32; np.log2 returning a non-integer triggers this ValueError.

Source

Thrown at jax/experimental/mosaic/gpu/utils.py:1810

  for s, t in zip(shape[-tiling_rank:], tiling):
    if s % t:
      raise ValueError("Non-divisible tiling:", shape, tiling)
  return (
      *shape[:-tiling_rank],
      *(s // t for s, t in zip(shape[-tiling_rank:], tiling)),
      *tiling,
  )


def warp_tree_reduce(value, op, group_size):
  """Reduce a value across the warpgroup."""
  assert bytewidth(value.type) == 4
  assert 32 % group_size == 0 and group_size <= 32
  i32 = ir.IntegerType.get_signless(32)
  result = value
  iters = np.log2(group_size)
  if not iters.is_integer():
    raise ValueError(
        f"Warp reduction group size should be a power of 2 (got {group_size})"
    )
  iters = int(iters)
  for i in range(iters):
    other_result = nvvm.shfl_sync(
        c(0xFFFFFFFF, i32),
        result,
        c(1 << i, i32),
        c(0x1F, i32),
        nvvm.ShflKind.bfly
    )
    result = op(result, other_result)

  return result


_MEMORY_SPACES = {f"#gpu.address_space<{str(x)}>": x for x in gpu.AddressSpace}

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Round group_size to the nearest power of two <= 32 (1, 2, 4, 8, 16, 32)
2. Check the computation of group_size (often threads_per_warp or an inferred worker count) upstream
3. If you need non-power-of-2 groups, pad the group to the next power of 2 and mask out inactive lanes

Example fix

# before
group_size = num_workers // 3  # e.g. 24 -> raises
# after
group_size = 1 << (num_workers // 3 - 1).bit_length()  # round up to power of 2

def valid_group_size(g):
    return g in (1, 2, 4, 8, 16, 32)
assert valid_group_size(group_size), f'group_size {group_size} not a power of 2 <= 32'

When it happens

Trigger: Calling the warp reduction helper with group_size values like 3, 6, 12, 24 — any non-power-of-2 value (the function also requires 32 % group_size == 0 and group_size <= 32, so effectively only 1,2,4,8,16,32 are valid).

Common situations: Deriving group_size from warp/thread configuration arithmetic (e.g. num_threads // something) that lands on 24 or 12; porting CUDA code that used arbitrary subgroup sizes; changing the number of lanes per worker in a Mosaic kernel.

Related errors

• packed cannot be specified if layout is specified.
• packed, collective and layout arguments are only supported f
• Cannot specify both out_shape and out_type. Use out_type.
• Cannot specify both scratch_shapes and scratch_types. Use sc
• Unsupported dtype: {ref.dtype}

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