{"record":{"id":"93a20c724037a254","repo":"jax-ml/jax","slug":"tmem-can-only-be-sliced-not-indexed","errorCode":null,"errorMessage":"TMEM can only be sliced, not indexed","messagePattern":"TMEM can only be sliced, not indexed","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/tcgen05.py","lineNumber":1248,"sourceCode":"    if shape[0] < 32:\n      raise ValueError(f\"TMEM refs must have at least 32 rows, got: {shape[0]}\")\n    if layout is None:\n      if collective is None:\n        raise ValueError(\n            \"collective argument must be provided when TMEM layout is inferred\"\n        )\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]):","sourceCodeStart":1230,"sourceCodeEnd":1266,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/tcgen05.py#L1230-L1266","documentation":"TMEM slicing produces a view, not a copy, so removing a dimension (indexing with an integer, which squeezes the axis) has no meaning for TMEM addressing. utils.parse_indices marks squeezed dims, and TMEMRef.slice rejects any non-None integer index.","triggerScenarios":"tmem_ref[0], tmem_ref[:, 3], or tmem_ref[5, :] — any integer index; also reached from apply_fun, broadcast_into, or lowering rules for subview/insert_strided_slice ops that index the TMEM ref.","commonSituations":"Writing numpy-style indexing on TMEM refs; MLIR-level subview lowering that fully reduces a dimension; converting register code that indexed accumulators.","solutions":["Use slices with explicit ranges: tmem_ref[0:shape0, a:b]","Select columns in registers after tcgen05.load instead of indexing TMEM","When lowering subviews, keep both dims (use 1-wide slices, but note row slicing limits)"],"exampleFix":"# before\nsub = tmem_ref[3]  # indexing\n# after\nsub = tmem_ref[:, 3:4]  # slicing (columns only; rows must stay full)","handlingStrategy":"validation","validationCode":"import numpy as np\nidxs = (slice(0, shape[0]), slice(a, b))\nassert all(isinstance(i, slice) for i in idxs), 'TMEM supports slicing only'","typeGuard":"def is_slice_only(idxs) -> bool:\n    return all(i is None or isinstance(i, slice) for i in idxs)","tryCatchPattern":null,"preventionTips":["Never index TMEM with integers; always use start:stop ranges","Select lanes/values in registers after load"],"tags":["jax","mosaic","tmem","slicing","unsupported-operation"],"backgroundTag":"unsupported-indexing","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}