{"record":{"id":"ddb4910289100eec","repo":"jax-ml/jax","slug":"cannot-slice-tmem-with-multiple-tiles-along-rows","errorCode":null,"errorMessage":"Cannot slice TMEM with multiple tiles along rows.","messagePattern":"Cannot slice TMEM with multiple tiles along rows\\.","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/tcgen05.py","lineNumber":1260,"sourceCode":"\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)\n\n    # The code below converts from a logical column index to a physical column\n    # index.\n    physical_cols_in_tile = self.layout.cols_in_shape(\n        cast(tuple[int, int], self.layout.base_tile_shape),\n        utils.bitwidth(self.dtype),","sourceCodeStart":1242,"sourceCodeEnd":1278,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/tcgen05.py#L1242-L1278","documentation":"After the row checks, slice() requires the ref's row count to equal the layout's base tile row count — with multiple row tiles, a column slice would produce non-contiguous memory across row tiles. Hence 'Cannot slice TMEM with multiple tiles along rows.' when self.shape[0] != layout.base_tile_shape[0].","triggerScenarios":"Calling slice (directly or via to_layout/apply_fun/subview lowering) on a TMEM ref whose layout tiles rows multiple times, e.g. a 128-row ref with a layout whose base tile covers 64 rows.","commonSituations":"Relayouting or column-slicing accumulators whose layout was constructed with smaller row tiles; mixing layouts from tmem_m64_collective_layout with full 128-row shapes.","solutions":["Avoid column slicing/relayout on multi-row-tile refs; operate on the full ref","Use a layout whose base tile rows equal the ref's rows (e.g. tmem_default_layout for 128 rows)","Load to registers first, then slice in the fragment-array domain"],"exampleFix":"# before\nsub = big_tmem_ref[:, 32:64]  # ref has multiple row tiles\n# after\nregs = tcgen05.load(big_tmem_ref)  # then slice registers, or use a single-row-tile layout","handlingStrategy":"validation","validationCode":"assert ref.shape[0] == ref.layout.base_tile_shape[0], 'multi-row-tile layout: cannot column-slice'","typeGuard":"def single_row_tile(ref) -> bool:\n    return ref.shape[0] == ref.layout.base_tile_shape[0]","tryCatchPattern":null,"preventionTips":["Use canonical layout constructors so base tile rows match the ref rows","Load to registers and slice there instead of slicing multi-tile TMEM refs"],"tags":["jax","mosaic","tmem","slicing","layout","unsupported-operation"],"backgroundTag":"unsupported-slicing","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}