{"record":{"id":"1d30969bf0d31727","repo":"jax-ml/jax","slug":"async-copies-only-support-striding-up-to-5-dimensi","errorCode":null,"errorMessage":"Async copies only support striding up to 5 dimensions","messagePattern":"Async copies only support striding up to 5 dimensions","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/launch_context.py","lineNumber":1125,"sourceCode":"      slice_shape = list(drop.transform_shape(slice_shape))\n      # After _prepare_async_copy, squeezed dims have been permuted to the\n      # front (via a TransposeTransform in gmem_transform). So in the\n      # transformed shape, they occupy the first `num_squeezed_dims`\n      # positions.\n      # TODO(bchetioui): move the creation of the `TransposeTransform`\n      # here instead of in _prepare_async_copy.\n      squeezed_dims = tuple(d for i, d in enumerate(squeezed_dims) if i not in unit_dims)\n      num_squeezed_dims = len(squeezed_dims)\n      if len(slice_shape) > 5 and squeezed_dims:\n        # We can try to collapse all squeezed dims into one.\n        squeezed_dim_strides = tuple(gmem_strides[d] for d in squeezed_dims)\n        collapse = CollapseLeadingIndicesTransform(squeezed_dim_strides)\n        gmem_transform = (*gmem_transform, collapse)\n        dyn_base_indices = collapse.transform_index(dyn_base_indices)\n        slice_shape = list(collapse.transform_shape(tuple(slice_shape)))\n        num_squeezed_dims = 1\n      if len(slice_shape) > 5:\n        raise ValueError(\"Async copies only support striding up to 5 dimensions\")\n    del squeezed_dims\n\n    # pyrefly: ignore[redefinition]\n    dyn_base_indices: list[ir.Value] = list(dyn_base_indices)\n    slice_shape = list(slice_shape)\n    assert all(d == 1 for d in slice_shape[:num_squeezed_dims])\n\n    # Partitioned loads have already been processed (before transforms).\n    # We process non-partitioned collective loads here, because only here are we\n    # able to know in what order the data will be written to SMEM. Transposes\n    # and tiling change that order and if we picked a partition based on the\n    # untransformed slice shape, we might have ended up with a non-contiguous\n    # SMEM window, which would no longer be realizable in a single TMA transfer.\n    collective_size = math.prod(self.cluster_size[d] for d in collective)\n    if collective_size > 1 and not isinstance(leader_tracked, _Partitioned):\n      assert gather_indices is None  # Checked above.\n      def partition_dim(dim: int, idx: ir.Value, num_chunks: int):\n        # No need to partition squeezed dims. They don't even exist in smem_ref.","sourceCodeStart":1107,"sourceCodeEnd":1143,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/launch_context.py#L1107-L1143","documentation":"TMA descriptors can describe at most 5 strided dimensions. After Mosaic squeezes/collapses singleton and leading index dimensions, if the slice still has more than 5 dimensions this ValueError is raised in _prepare_tma. The library mirrors the hardware limit of the TMA tensor map.","triggerScenarios":"Calling async_copy/async_prefetch with the TMA implementation on a slice whose shape, after collapsing leading unit-stride dimensions, still has rank > 5, e.g. an effective 6D tile.","commonSituations":"Operating on high-rank tensors (e.g. batched attention with heads/batch/seq/heads dims) in Mosaic GPU kernels; adding dimensions via transforms instead of reshaping first.","solutions":["Reshape the tensor or slice so the copy spans at most 5 dimensions before calling async_copy.","Ensure leading dimensions can be collapsed (make them contiguous/unit-stride) so the CollapseLeadingIndicesTransform reduces the rank below 5.","Split the copy into multiple async_copy calls over lower-rank slices."],"exampleFix":"// before\nctx.async_copy(gmem_ref, smem_ref, gmem_slice=tuple_of_6d_slices)\n// after\nx = x.reshape(x.shape[:2] + (-1,) + x.shape[4:])  # fold dims to <=5D\nctx.async_copy(gmem_ref, smem_ref, gmem_slice=fewer_dim_slices)","handlingStrategy":"validation","validationCode":"effective_rank = len([d for d in slice_shape if d > 1])  # after collapsing unit dims\nassert effective_rank <= 5 or len(slice_shape) <= 5, 'TMA supports at most 5 strided dims'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Reshape operands to <=5D before entering Mosaic kernels.","Keep layouts contiguous so leading dims collapse.","Split high-rank copies into multiple lower-rank async_copy calls."],"tags":["jax","mosaic-gpu","tma","rank-limit","shape-validation"],"backgroundTag":"tensor-rank-limit","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}