jax-ml/jax · error · ValueError

None of the leading dimensions in the transformed slice shap

Error message

None of the leading dimensions in the transformed slice shape {slice_shape} is divisible by the collective size {collective_size}

What it means

When using TMA with collective_size > 1, Mosaic tries to partition the copy across leading dimensions of the (transformed) slice shape by dividing dimensions by the collective size. If none of the leading dimensions is divisible, partitioning fails and this ValueError is raised.

Source

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

      for dim, slice_size in enumerate(
          slice_shape[:-1] if has_swizzle else slice_shape
      ):
        if slice_size % rem_collective_size == 0:
          partition_dim(dim, idx, rem_collective_size)
          rem_collective_size = 1
          break
        elif rem_collective_size % slice_size == 0:
          # This is an optimization and it lets us skip squeezed dims.
          if slice_size > 1:
            dim_idx = arith.remui(idx, c(slice_size, index))
            partition_dim(dim, dim_idx, slice_size)
            idx = arith.divui(idx, c(slice_size, index))
            rem_collective_size //= slice_size
        else:
          break  # We failed to partition the leading dimensions.
      del idx  # We overwrote the block index in the loop.
      if rem_collective_size > 1:
        raise ValueError(
            "None of the leading dimensions in the transformed slice shape"
            f" {slice_shape} is divisible by the collective size"
            f" {collective_size}"
        )

    if (zeroth_bw := slice_shape[-1] * element_bitwidth) % 128 != 0:
      raise ValueError(
          "Async copies require the number of bits copied along the last"
          f" dimension to be divisible by 128, but got {zeroth_bw}"
      )
    if (
        swizzle is not None
        and swizzle != mgpu_dialect.SwizzlingMode.kNoSwizzle
        and slice_shape[-1] != (swizzle * 8) // element_bitwidth
    ):
      raise ValueError(
          f"Async copies with {swizzle=} require the last dimension of the"
          f" slice to be exactly {swizzle} bytes i.e. "

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Pad or resize at least one leading dimension of the slice so it is divisible by the collective size.
2. Reshape so the dimension carrying parallelism (e.g. batch or sequence) is the leading dim and is a multiple of collective_size.
3. Reduce collective_size to 1 if partitioning across CTAs is not required.

Example fix

// before
ctx.async_copy(..., collective=(2,), gmem_slice=(slice(0,3), slice(0,64)))
// after
ctx.async_copy(..., collective=(2,), gmem_slice=(slice(0,4), slice(0,64)))

if collective_size > 1:
    assert any(d % collective_size == 0 for d in slice_shape[:-1]), \
        'a leading dim must be divisible by collective_size for TMA partitioning'

try:
    ctx.async_copy(..., collective=(2,))
except ValueError as e:
    if 'None of the leading dimensions' in str(e):
        slice_shape[0] = _round_up(slice_shape[0], collective_size)
    else:
        raise

When it happens

Trigger: Calling async_copy/async_prefetch with implementation=TMA and a collective size > 1 where no leading dimension of the transformed slice shape is divisible by the collective size (e.g. all leading dims are 1 or odd sizes with collective size 2).

Common situations: Multicast/partitioned TMA loads on Hopper clusters where the tensor layout wasn't padded for cluster partitioning; changing collective configuration without adjusting tile shapes.

Related errors

• The collective size ({collective_size}) must divide the slic
• Async copies only support striding up to 5 dimensions
• Async copies require the number of bits copied along the las
• Only byte-aligned shapes are supported. Got shape: {ref.dtyp
• copy_gmem_to_smem with a barrier is only supported Hopper an

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