{"record":{"id":"2f266e48ce6457b6","repo":"jax-ml/jax","slug":"tmem-reference-must-have-tmem-rows-rows-but-got","errorCode":null,"errorMessage":"TMEM reference must have {TMEM_ROWS} rows, but got {tmem_ref.shape[0]}","messagePattern":"TMEM reference must have (.+?) rows, but got (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/tcgen05.py","lineNumber":2014,"sourceCode":"    tmem_ref: TMEMRef,\n    swizzle: int,\n    collective: bool = False,\n) -> None:\n  i8 = ir.IntegerType.get_signless(8)\n  i32 = ir.IntegerType.get_signless(32)\n  smem_ty = ir.MemRefType(smem_ref.type)\n  if (dtype := smem_ty.element_type) != tmem_ref.dtype:\n    raise ValueError(f\"Incompatible dtypes: SMEM has {dtype}, TMEM has {tmem_ref.dtype}\")\n  if swizzle not in {16, 32, 64, 128}:\n    raise ValueError(f\"Unsupported swizzle, expected 16, 32, 64 or 128, but got: {swizzle}\")\n  bitwidth = utils.bitwidth(dtype)\n  if tmem_ref.packing != 32 // bitwidth:\n    raise ValueError(\n        \"tcgen05.cp only supports fully packed TMEM references\"\n        f\" (packing={32 // bitwidth}), but got packing={tmem_ref.packing}\"\n    )\n  if tmem_ref.shape[0] != TMEM_ROWS:\n    raise ValueError(\n        f\"TMEM reference must have {TMEM_ROWS} rows, but got {tmem_ref.shape[0]}\"\n    )\n  if tmem_ref.layout != tmem_default_layout(packing=tmem_ref.packing):\n    raise ValueError(\n        f\"Only standard TMEM layout is supported, got: {tmem_ref.layout}\"\n    )\n  swizzle_elems = 8 * swizzle // bitwidth\n  expected_smem_shape = utils.tile_shape(tmem_ref.shape, (8, swizzle_elems))\n  smem_shape = tuple(smem_ty.shape)\n  if smem_shape != expected_smem_shape:\n    raise ValueError(\n        f\"SMEM has shape {smem_shape}, but expected {expected_smem_shape} for\"\n        f\" TMEM shape {tmem_ref.shape} with swizzle={swizzle}\"\n    )\n  strides, _ = smem_ty.get_strides_and_offset()\n  row_tile_stride, col_tile_stride, inner_row_stride, inner_col_stride = strides\n  if inner_col_stride != 1 or inner_row_stride != swizzle_elems:\n    raise ValueError(\"The SMEM tiles must be contiguous\")","sourceCodeStart":1996,"sourceCodeEnd":2032,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/tcgen05.py#L1996-L2032","documentation":"Raised by async_copy_smem_to_tmem when the TMEM reference's row count (shape[0]) does not equal TMEM_ROWS (128). The tcgen05.cp instruction writes a full 128-row TMEM block per copy, so the destination reference must span all 128 lanes.","triggerScenarios":"Passing a tmem_ref that was sliced to fewer than 128 rows (e.g. tmem_ref.slice or a per-warpgroup 64-row view), or allocating TMEM with an M dimension other than 128, to async_copy_smem_to_tmem.","commonSituations":"Splitting TMEM across warpgroups or warps (each getting 32/64 rows) and then attempting a bulk SMEM->TMEM copy on the sub-view; using a matmul_shape with M != 128.","solutions":["Allocate/copy into a full 128-row TMEM reference, then slice per-warpgroup afterwards","Check the TMEM_ROWS constant (128) against your tmem_alloc shape[0]","Redesign the kernel so the SMEM staging tile covers all 128 rows in one copy"],"exampleFix":"// before\ntmem = tmem_alloc(dtype, (64, n))\nasync_copy_smem_to_tmem(smem, tmem, swizzle=32)\n// after\ntmem = tmem_alloc(dtype, (128, n))\nasync_copy_smem_to_tmem(smem, tmem, swizzle=32)","handlingStrategy":"validation","validationCode":"assert tmem_ref.shape[0] == 128, 'tcgen05.cp requires full 128-row TMEM'","typeGuard":"def is_full_tmem(tmem_ref, rows=128):\n    return tmem_ref.shape[0] == rows","tryCatchPattern":null,"preventionTips":["Do bulk SMEM->TMEM copies on the full 128-row reference before slicing per warpgroup","Check matmul M dimension equals 128 when planning tcgen05.cp usage"],"tags":["gpu","mosaic","tcgen05","tensor-memory","shape-mismatch"],"backgroundTag":"shape-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}