jax-ml/jax · error · ValueError

The collective size ({collective_size}) must divide the slic

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.

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)

Solutions

1. 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).
2. Verify math.prod(launch_context.cluster_size) == 2 and use a partitioned dimension whose size is a multiple of collective_size.
3. 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)

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:
    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

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

• None of the leading dimensions in the transformed slice shap
• Only byte-aligned shapes are supported. Got shape: {ref.dtyp
• Async copies only support striding up to 5 dimensions
• Async copies require the number of bits copied along the las
• Expected same element type, got {element_type} and {dst_ref_

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