{"record":{"id":"b7c5f0e2013e92f3","repo":"jax-ml/jax","slug":"tmem-cannot-be-sliced-along-rows","errorCode":null,"errorMessage":"TMEM cannot be sliced along rows","messagePattern":"TMEM cannot be sliced along rows","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/tcgen05.py","lineNumber":1254,"sourceCode":"        )\n      layout = _infer_tmem_layout(shape, collective, packing=1)\n    # TODO: Do we have to do this??\n    # warp_idx = utils.warp_idx(sync=False)\n    # tmem_addr = arith.ori(tmem_addr, arith.shli(warp_idx, utils.c(21, i32)))\n    return cls(tmem_addr, shape, dtype, layout)\n\n  def slice(self, *idxs) -> TMEMRef:\n    i32 = ir.IntegerType.get_signless(32)\n    base_idx, slice_shape, is_squeezed = utils.parse_indices(idxs, self.shape)\n    slice_shape = cast(tuple[int, int], tuple(slice_shape))\n    if any(is_squeezed):\n      raise ValueError(\"TMEM can only be sliced, not indexed\")\n    if base_idx == [0] * len(base_idx) and slice_shape == self.shape:\n      return self  # Trivial slice\n    # If we slice along rows, or attempt to extract several rows, then we may\n    # end up with a non-contiguous slice of memory.\n    if base_idx[0] != 0 or slice_shape[0] != self.shape[0]:\n      raise NotImplementedError(\"TMEM cannot be sliced along rows\")\n    # If we attempt to extract non-contiguous tiles, then we will end up with a\n    # non-contiguous slice of memory.\n    # We check that we have a single tile along rows. Hence slicing along\n    # columns produces a contiguous slice of memory.\n    if self.shape[0] != self.layout.base_tile_shape[0]:\n      raise NotImplementedError(\n          \"Cannot slice TMEM with multiple tiles along rows.\"\n      )\n    col_idx = base_idx[1]\n    if not isinstance(col_idx, ir.Value):\n      col_idx = arith.constant(i32, col_idx)\n    if not utils.is_known_divisible(col_idx, self.layout.base_tile_shape[1]):\n      raise NotImplementedError(\n          \"Slicing along columns is not supported when the column index is not\"\n          \" known to be a multiple of the base tile shape\"\n      )\n    if col_idx.type == ir.IndexType.get():\n      col_idx = arith.index_cast(i32, col_idx)","sourceCodeStart":1236,"sourceCodeEnd":1272,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/tcgen05.py#L1236-L1272","documentation":"TMEM rows map to lane/DPMC structure that is not contiguous for arbitrary row ranges, so a slice must start at row 0 and keep all rows. slice() raises NotImplementedError when base_idx[0] != 0 or slice_shape[0] != full row count.","triggerScenarios":"tmem_ref[10:, :], tmem_ref[0:64, 8:16] on a 128-row ref, or subview lowerings that shrink the row extent.","commonSituations":"Splitting an accumulator across warps/CTAs by row; adapting SMEM subview code to TMEM; trying to feed half of an M=128 accumulator into two MMAs.","solutions":["Slice only along columns: tmem_ref[:, a:b]","Allocate separate 64-row TMEM refs instead of slicing a 128-row one by rows","Move row partitioning into the MMA shape or register-level operations"],"exampleFix":"# before\nhalf = tmem_ref[0:64, :]  # row slicing on 128-row ref\n# after\nhalf = tcgen05.TMEMRef.from_alloc(alloc64, (64, ncols), collective=True)  # separate 64-row alloc","handlingStrategy":"validation","validationCode":"assert r0 == 0 and r1 == ref.shape[0], 'TMEM row slicing unsupported'","typeGuard":null,"tryCatchPattern":"try:\n    sub = ref[idx]\nexcept NotImplementedError as e:\n    if 'along rows' in str(e):\n        raise  # restructure: separate allocs, not row slices\n    raise","preventionTips":["Design kernels to slice TMEM by columns only","Partition rows via separate allocations or MMA shapes, never via slices"],"tags":["jax","mosaic","tmem","slicing","unsupported-operation"],"backgroundTag":"unsupported-slicing","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}