{"record":{"id":"fb11e69341e22f2d","repo":"jax-ml/jax","slug":"tmem-refs-must-have-at-least-32-rows-got-shape","errorCode":null,"errorMessage":"TMEM refs must have at least 32 rows, got: {shape[0]}","messagePattern":"TMEM refs must have at least 32 rows, got: (.+?)","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/tcgen05.py","lineNumber":1231,"sourceCode":"      shape: tuple[int, int],\n      dtype,\n      collective: bool | None = None,\n      layout: TMEMLayout | None = None,\n  ) -> TMEMRef:\n    i32 = ir.IntegerType.get_signless(32)\n    if not isinstance(tmem_addr_ref.type, ir.MemRefType):\n      raise ValueError(f\"tmem_addr_ref must be a memref or a pointer, got: {tmem_addr_ref.type}\")\n    addr_ref_ty = ir.MemRefType(tmem_addr_ref.type)\n    if not utils.is_smem_ref(addr_ref_ty):\n      raise ValueError(f\"tmem_addr_ref must be in shared memory, got: {addr_ref_ty}\")\n    if addr_ref_ty.element_type != i32:\n      raise ValueError(f\"tmem_addr_ref must be an i32 memref, got: {addr_ref_ty}\")\n    if math.prod(addr_ref_ty.shape) != 1:\n      raise ValueError(f\"tmem_addr_ref must contain a single element, got: {addr_ref_ty}\")\n    i0 = arith.ConstantOp.create_index(0)\n    tmem_addr = memref.load(tmem_addr_ref, [i0] * addr_ref_ty.rank)\n    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:","sourceCodeStart":1213,"sourceCodeEnd":1249,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/tcgen05.py#L1213-L1249","documentation":"TMEM is allocated in units of 32 lanes minimum on Blackwell; from_alloc rejects shapes with fewer than 32 rows even though the layout inference only supports 64/128 rows. This is an early guard against degenerate allocations that the hardware cannot make.","triggerScenarios":"TMEMRef.from_alloc(alloc, (16, N), ...) or any shape[0] < 32.","commonSituations":"Trying small test shapes like (8, 64) or (16, 16) when prototyping tcgen05 kernels; slicing a TMEM ref down below 32 rows (slice itself forbids row slicing, so this mostly occurs at allocation time).","solutions":["Use at least 32 rows — practically 64 or 128 to satisfy layout inference","Keep small matrices in registers/SMEM instead of TMEM","Scale test shapes up to the hardware minimum"],"exampleFix":"# before\nref = tcgen05.TMEMRef.from_alloc(alloc, (16, 64), collective=True)\n# after\nref = tcgen05.TMEMRef.from_alloc(alloc, (64, 64), collective=True)","handlingStrategy":"validation","validationCode":"assert shape[0] >= 32, shape","typeGuard":"def tmem_rows_valid(shape) -> bool:\n    return len(shape) == 2 and shape[0] >= 32","tryCatchPattern":null,"preventionTips":["Don't prototype with tiny TMEM shapes; hardware minimum is 32 lanes and layouts need 64/128","Keep small operands in registers or SMEM"],"tags":["jax","mosaic","tmem","shape-validation","hardware-constraint"],"backgroundTag":"shape-validation-failed","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}