{"record":{"id":"c6a6c6811f27c04c","repo":"jax-ml/jax","slug":"the-collective-size-collective-size-must-divid","errorCode":null,"errorMessage":"The collective size ({collective_size}) must divide the slice shape along the partitioned dimension, but it has size {slice_shape[partitioned]}","messagePattern":"The collective size \\((.+?)\\) must divide the slice shape along the partitioned dimension, but it has size (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/launch_context.py","lineNumber":1008,"sourceCode":"      if gather_indices is not None:\n        raise NotImplementedError(\"Collective copies with gather/scatter unsupported\")\n    if isinstance(leader_tracked, _Partitioned):\n      partitioned = leader_tracked.axis\n      # Increment partitioned by the number of preceding squeezed dimensions.\n      partitioned = np.where(\n          np.cumsum(~np.array(is_squeezed)) == partitioned+1)[0][0]\n      # Partitioning happens on the logical slice we extract from GMEM, so we do\n      # it before we apply transforms.\n      if not collective:  # This implies non-gather TMA already.\n        raise ValueError(\"Only collective loads can be partitioned\")\n      collective_size = math.prod(self.cluster_size[d] for d in collective)\n      if collective_size > 1:\n        if math.prod(self.cluster_size) != 2:\n          raise NotImplementedError(\n              \"Partitioned loads only supported for clusters of size 2\"\n          )\n        if slice_shape[partitioned] % collective_size != 0:\n          raise ValueError(\n              f\"The collective size ({collective_size}) must divide the slice\"\n              \" shape along the partitioned dimension, but it has size\"\n              f\" {slice_shape[partitioned]}\"\n          )\n        slice_shape[partitioned] //= collective_size\n        dyn_base_indices = list(dyn_base_indices)\n        dyn_base_indices[partitioned] = arith.addi(\n            dyn_base_indices[partitioned],\n            arith.muli(\n                utils.cluster_idx(collective),\n                c(slice_shape[partitioned], index),\n            ),\n        )\n        dyn_base_indices = tuple(dyn_base_indices)\n\n    squeezed_dims = tuple(\n        i for i, squeezed in enumerate(is_squeezed) if squeezed\n    )","sourceCodeStart":990,"sourceCodeEnd":1026,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/launch_context.py#L990-L1026","documentation":"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.","triggerScenarios":"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.","commonSituations":"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).","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."],"exampleFix":"// before\nctx.async_copy(gmem_ref, smem_ref, gmem_slice=(slice(0, 5),), collective=coll)\n// after\n# pad/align the slice so the partitioned dim is divisible by the cluster size\nctx.async_copy(gmem_ref, smem_ref, gmem_slice=(slice(0, 6),), collective=coll)","handlingStrategy":"validation","validationCode":"prod = math.prod(cluster_size)\nassert prod in (1, 2), 'partitioned loads only support clusters of size 2'\nif prod > 1 and slice_shape[partitioned_dim] % collective_size != 0:\n    raise ValueError(f'pad slice dim {partitioned_dim} to a multiple of {collective_size}')","typeGuard":null,"tryCatchPattern":"try:\n    ctx.async_copy(...)\nexcept ValueError as e:\n    if 'must divide the slice shape' in str(e):\n        slice_shape[partitioned_dim] = _round_up(slice_shape[partitioned_dim], collective_size)\n    else:\n        raise","preventionTips":["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."],"tags":["jax","mosaic-gpu","async-copy","cluster-partitioning","shape-validation"],"backgroundTag":"shape-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}