{"record":{"id":"7f32b7e97b200794","repo":"jax-ml/jax","slug":"none-of-the-leading-dimensions-in-the-transformed","errorCode":null,"errorMessage":"None of the leading dimensions in the transformed slice shape {slice_shape} is divisible by the collective size {collective_size}","messagePattern":"None of the leading dimensions in the transformed slice shape (.+?) is divisible by the collective size (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/launch_context.py","lineNumber":1180,"sourceCode":"      for dim, slice_size in enumerate(\n          slice_shape[:-1] if has_swizzle else slice_shape\n      ):\n        if slice_size % rem_collective_size == 0:\n          partition_dim(dim, idx, rem_collective_size)\n          rem_collective_size = 1\n          break\n        elif rem_collective_size % slice_size == 0:\n          # This is an optimization and it lets us skip squeezed dims.\n          if slice_size > 1:\n            dim_idx = arith.remui(idx, c(slice_size, index))\n            partition_dim(dim, dim_idx, slice_size)\n            idx = arith.divui(idx, c(slice_size, index))\n            rem_collective_size //= slice_size\n        else:\n          break  # We failed to partition the leading dimensions.\n      del idx  # We overwrote the block index in the loop.\n      if rem_collective_size > 1:\n        raise ValueError(\n            \"None of the leading dimensions in the transformed slice shape\"\n            f\" {slice_shape} is divisible by the collective size\"\n            f\" {collective_size}\"\n        )\n\n    if (zeroth_bw := slice_shape[-1] * element_bitwidth) % 128 != 0:\n      raise ValueError(\n          \"Async copies require the number of bits copied along the last\"\n          f\" dimension to be divisible by 128, but got {zeroth_bw}\"\n      )\n    if (\n        swizzle is not None\n        and swizzle != mgpu_dialect.SwizzlingMode.kNoSwizzle\n        and slice_shape[-1] != (swizzle * 8) // element_bitwidth\n    ):\n      raise ValueError(\n          f\"Async copies with {swizzle=} require the last dimension of the\"\n          f\" slice to be exactly {swizzle} bytes i.e. \"","sourceCodeStart":1162,"sourceCodeEnd":1198,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/launch_context.py#L1162-L1198","documentation":"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.","triggerScenarios":"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).","commonSituations":"Multicast/partitioned TMA loads on Hopper clusters where the tensor layout wasn't padded for cluster partitioning; changing collective configuration without adjusting tile shapes.","solutions":["Pad or resize at least one leading dimension of the slice so it is divisible by the collective size.","Reshape so the dimension carrying parallelism (e.g. batch or sequence) is the leading dim and is a multiple of collective_size.","Reduce collective_size to 1 if partitioning across CTAs is not required."],"exampleFix":"// before\nctx.async_copy(..., collective=(2,), gmem_slice=(slice(0,3), slice(0,64)))\n// after\nctx.async_copy(..., collective=(2,), gmem_slice=(slice(0,4), slice(0,64)))","handlingStrategy":"validation","validationCode":"if collective_size > 1:\n    assert any(d % collective_size == 0 for d in slice_shape[:-1]), \\\n        'a leading dim must be divisible by collective_size for TMA partitioning'","typeGuard":null,"tryCatchPattern":"try:\n    ctx.async_copy(..., collective=(2,))\nexcept ValueError as e:\n    if 'None of the leading dimensions' in str(e):\n        slice_shape[0] = _round_up(slice_shape[0], collective_size)\n    else:\n        raise","preventionTips":["Pad leading dims to multiples of the collective size.","Put the partitioned/parallel dimension first.","Test with the exact cluster shape you launch with."],"tags":["jax","mosaic-gpu","tma","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"}