{"record":{"id":"f25a2cb71d580f61","repo":"jax-ml/jax","slug":"unsupported-shape-shape-tmem-references-must-h","errorCode":null,"errorMessage":"Unsupported shape: {shape}. TMEM references must have either {TMEM_ROWS} or {TMEM_ROWS // 2} rows, but got: {shape[0]}.","messagePattern":"Unsupported shape: (.+?)\\. TMEM references must have either (.+?) or (.+?) rows, but got: (.+?)\\.","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/tcgen05.py","lineNumber":1078,"sourceCode":"    )\n\n\ndef _infer_tmem_layout(shape: tuple[int, ...], collective: bool, packing: int) -> TMEMLayout:\n  if len(shape) != 2:\n    raise ValueError(f\"TMEM can only represent 2D shapes, got {shape}\")\n  if packing > 8 or packing.bit_count() != 1:\n    raise ValueError(f\"Packing must be <= 8 and a power of 2, got: {packing}\")\n  if shape[1] % packing:\n    raise ValueError(f\"Minor dimension of shape must be divisible by packing, got: {shape}\")\n  if shape[0] == TMEM_ROWS:\n    return tmem_default_layout(packing)\n  elif shape[0] == TMEM_ROWS // 2:\n    if collective:\n      return tmem_m64_collective_layout(shape[1], packing)\n    else:\n      return tmem_half_lane_layout(shape[1], packing)\n  else:\n    raise ValueError(\n        f\"Unsupported shape: {shape}. TMEM references must have either\"\n        f\" {TMEM_ROWS} or {TMEM_ROWS // 2} rows, but got {shape[0]}.\"\n    )\n\n\ndef tmem_default_layout(packing: int = 1) -> TMEMLayout:\n  \"\"\"A TMEM layout used for 1CTA MMA with M=128 and 2CTA MMA with M=256.\"\"\"\n  if packing.bit_count() != 1:\n    raise ValueError(f\"Packing must be a power of 2, got: {packing}\")\n  return TMEMLayout(\n      fa.Tiling(((TMEM_ROWS, packing), (fa.WARP_SIZE, packing))),\n      warp_dims=(-4,),\n      lane_dims=(-2,),\n      vector_dim=-1,\n  )\n\n\ndef tmem_half_lane_layout(columns, packing: int = 1) -> TMEMLayout:","sourceCodeStart":1060,"sourceCodeEnd":1096,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/tcgen05.py#L1060-L1096","documentation":"TMEM on Blackwell is organized into 128 lanes (rows); hardware only supports allocations of 128 rows or 64 rows (half-lane). _infer_tmem_layout dispatches on shape[0] being TMEM_ROWS (128) or TMEM_ROWS//2 (64) and rejects any other row count because no TMEM layout can describe it.","triggerScenarios":"from_alloc / infer_tmem_layout / mma with an accumulator shape like (32, N), (96, N), (256, N) — any row count other than 128 or 64. Note from_alloc separately requires >=32 rows, so values like 32 or 96 reach this error.","commonSituations":"Porting WGMMA (Hopper) kernels that used M=32/48/96 to tcgen05; computing TMEM shape from arbitrary MMA M dimensions; multi-CTA code where the per-CTA slice of M is not 64 or 128.","solutions":["Reshape/split the computation so the TMEM-resident operand has 128 or 64 rows (e.g. M=256 2CTA -> 128 rows per CTA with collective=True)","For small M, keep registers in SMEM/registers rather than TMEM","Pass an explicit valid layout with a compatible shape instead of inference"],"exampleFix":"# before\nref = tcgen05.TMEMRef.from_alloc(alloc, (96, 64), collective=True)\n# after\nref = tcgen05.TMEMRef.from_alloc(alloc, (128, 64), collective=True)  # or (64, 64)","handlingStrategy":"validation","validationCode":"TMEM_ROWS = 128\nassert shape[0] in (TMEM_ROWS, TMEM_ROWS // 2), f'bad TMEM rows: {shape[0]}'","typeGuard":"def is_supported_tmem_rows(shape: tuple[int, ...]) -> bool:\n    return len(shape) == 2 and shape[0] in (128, 64)","tryCatchPattern":null,"preventionTips":["Constrain MMA M to 64/128 (or 128/256 for 2CTA) at kernel design time","Add a unit assert on accumulator shape before from_alloc/mma"],"tags":["jax","mosaic","tmem","shape-validation","blackwell"],"backgroundTag":"shape-validation-failed","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}