jax-ml/jax · error · ValueError

The collective size ( ) must divide the slice shape along…

Error message

The collective size ({collective_size}) must divide the slice shape along the partitioned dimension, but it has size {slice_shape[partitioned]}

What it means

Mosaic GPU's async_copy partitions a slice across the collective (cluster) when collective_size > 1. For the non-TMA partitioned-load path, only clusters of size 2 are supported, and the slice shape along the partitioned dimension must be evenly divisible by the collective size. If slice_shape[partitioned] % collective_size != 0, this ValueError is raised from _prepare_async_copy.

Solutions

Make the partitioned dimension extent divisible by the collective size (pad or reshape the slice, e.g. slice an even number of elements for a 2-CTA cluster).
Verify math.prod(launch_context.cluster_size) == 2 and use a partitioned dimension whose size is a multiple of collective_size.
Switch to the TMA implementation (AsyncCopyImplementation.TMA), whose partitioning logic differs and supports divisibility across leading dimensions.

Example fix

// before
ctx.async_copy(gmem_ref, smem_ref, gmem_slice=(slice(0, 5),), collective=coll)
// after
# pad/align the slice so the partitioned dim is divisible by the cluster size
ctx.async_copy(gmem_ref, smem_ref, gmem_slice=(slice(0, 6),), collective=coll)

Defensive patterns

Strategy: validation

Validate before calling

prod = math.prod(cluster_size)
assert prod in (1, 2), 'partitioned loads only support clusters of size 2'
if prod > 1 and slice_shape[partitioned_dim] % collective_size != 0:
    raise ValueError(f'pad slice dim {partitioned_dim} to a multiple of {collective_size}')

Try / catch

try:
    ctx.async_copy(...)
except ValueError as e:
    if 'must divide the slice shape' in str(e):
        slice_shape[partitioned_dim] = _round_up(slice_shape[partitioned_dim], collective_size)
    else:
        raise

Prevention

Keep partitioned-dimension tile sizes powers of two.
Assert cluster size is 2 before using partitioned non-TMA loads.
Write unit tests over multiple slice shapes for cluster kernels.

When it happens

Trigger: Calling async_copy or async_prefetch with a collective/cluster size > 1 (cluster_size with prod == 2) where the slice extent along the partitioned dimension is not divisible by the collective size, e.g. slicing 5 elements with collective size 2.

Common situations: Writing Mosaic GPU kernels that use TMA-less partitioned loads on Hopper/Blackwell clusters; changing tensor shapes or slice sizes without updating the cluster partitioning; assuming arbitrary collective sizes are supported (only size 2 is).

Related errors

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

Appendix: source

Thrown at jax/experimental/mosaic/gpu/launch_context.py:1008

      if gather_indices is not None:
        raise NotImplementedError("Collective copies with gather/scatter unsupported")
    if isinstance(leader_tracked, _Partitioned):
      partitioned = leader_tracked.axis
      # Increment partitioned by the number of preceding squeezed dimensions.
      partitioned = np.where(
          np.cumsum(~np.array(is_squeezed)) == partitioned+1)[0][0]
      # Partitioning happens on the logical slice we extract from GMEM, so we do
      # it before we apply transforms.
      if not collective:  # This implies non-gather TMA already.
        raise ValueError("Only collective loads can be partitioned")
      collective_size = math.prod(self.cluster_size[d] for d in collective)
      if collective_size > 1:
        if math.prod(self.cluster_size) != 2:
          raise NotImplementedError(
              "Partitioned loads only supported for clusters of size 2"
          )
        if slice_shape[partitioned] % collective_size != 0:
          raise ValueError(
              f"The collective size ({collective_size}) must divide the slice"
              " shape along the partitioned dimension, but it has size"
              f" {slice_shape[partitioned]}"
          )
        slice_shape[partitioned] //= collective_size
        dyn_base_indices = list(dyn_base_indices)
        dyn_base_indices[partitioned] = arith.addi(
            dyn_base_indices[partitioned],
            arith.muli(
                utils.cluster_idx(collective),
                c(slice_shape[partitioned], index),
            ),
        )
        dyn_base_indices = tuple(dyn_base_indices)

    squeezed_dims = tuple(
        i for i, squeezed in enumerate(is_squeezed) if squeezed
    )

View on GitHub (pinned to 1e1c6a8fc0)