jax-ml/jax · error · ValueError

Expected exactly one collective axis, got {collective_axes=}

Error message

Expected exactly one collective axis, got {collective_axes=}

What it means

For leader-tracked collective copies, both CTAs in the cluster send completion messages to the leader, so the code requires exactly one collective axis (len(collective) == 1); otherwise the accounting of doubled bytes is ambiguous and ValueError is raised.

Source

Thrown at jax/_src/pallas/mosaic_gpu/primitives.py:985

        lowering._resolve_cluster_axis(ctx.module_ctx.axis_names, axis)
        for axis in collective_axes
    )

  is_leader_tracked_copy = collective and leader_tracked is not None
  dst_ty = ir.MemRefType(dst.type)
  bits = math.prod(dst_ty.shape) * mgpu.bitwidth(dst_ty.element_type)
  if bits % 8:
    raise ValueError(
        f"Can only transfer integer bytes (shape={dst_ty.shape},"
        f" dtype={dst_ty.element_type})"
    )
  bytes = bits // 8

  if is_leader_tracked_copy:
    # Leader receives the completion messages from both CTAs.
    bytes *= 2
    if len(collective) != 1:
      raise ValueError(
          f"Expected exactly one collective axis, got {collective_axes=}"
      )
    if math.prod(ctx.launch_ctx.cluster_size) != 2:
      raise NotImplementedError(
          "Partitioned loads only supported for clusters of size 2. Got"
          f" cluster size {ctx.launch_ctx.cluster_size}."
      )

  # TMA is only available on Hopper and newer. On older architectures we fall
  # back to the cp.async implementation.
  if is_cp_async := mgpu.utils.get_arch().major < 9:
    if barrier is not None:
      raise ValueError(
          "copy_gmem_to_smem with a barrier is only supported Hopper and newer"
          " GPUs, which use the TMA implementation"
      )
    if collective_axes is not None:
      raise ValueError("Only the TMA implementation supports collective copies")

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Pass exactly one collective axis name, e.g. collective_axes=('data',)
2. Drop leader_tracked if you don't need leader completion tracking
3. Restructure the copy into separate per-axis collective copies

Example fix

# before
copy_gmem_to_smem(src, smem, collective_axes=('data','model'), leader_tracked=True)
# after
copy_gmem_to_smem(src, smem, collective_axes=('data',), leader_tracked=True)

if leader_tracked is not None:
    assert len(collective_axes) == 1, 'leader_tracked requires exactly one collective axis'

When it happens

Trigger: copy_gmem_to_smem with leader_tracked set and collective_axes containing zero or more than one axis names.

Common situations: Reusing a multi-axis mesh tuple for collective_axes when the design only partitions over one axis; passing collective_axes=() with leader_tracked still enabled.

Related errors

• collective_axes is not supported in pallas_call. Use plgpu.k
• Only collective_axes that include all JAX device mesh axes
• Partitioned loads only supported for clusters of size 2. Got
• Only the TMA implementation supports collective copies
• Only the TMA implementation supports leader_tracked copies

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