{"record":{"id":"e770acc61a0031af","repo":"jax-ml/jax","slug":"packing-must-be-a-power-of-2-got-packing","errorCode":null,"errorMessage":"Packing must be a power of 2, got: {packing}","messagePattern":"Packing must be a power of 2, got: (.+?)","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/tcgen05.py","lineNumber":1087,"sourceCode":"    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:\n  \"\"\"A TMEM layout used for 1CTA MMA with M=64.\"\"\"\n  if packing > (columns // 2) or packing.bit_count() != 1:\n    raise ValueError(f\"Packing must be <= 8 and a power of 2, got: {packing}\")\n  if columns % 16:\n    raise ValueError(f\"Columns must be a multiple of 16, got: {columns}\")\n  return TMEMLayout(\n      fa.Tiling((\n          (TMEM_ROWS // 2, columns),\n          (fa.WARP_SIZE // 2, columns // 2),","sourceCodeStart":1069,"sourceCodeEnd":1105,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/tcgen05.py#L1069-L1105","documentation":"tmem_default_layout builds the standard 128-row TMEM layout with a packing (vector width along columns) that must be a power of 2 because it maps directly onto vectorized TMEM loads/stores. packing.bit_count() != 1 catches zero, negatives, and non-powers like 3 or 6.","triggerScenarios":"Calling tmem_default_layout(3), tmem_default_layout(0), or passing a computed packing such as bitwidth-derived values that aren't powers of two; also reached indirectly via layout-related helpers (is_valid_tmem_transfer, pprint_layout, async store constraint system).","commonSituations":"Deriving packing from element bitwidth with arithmetic that yields non-power-of-2 (e.g. 24/8=3 is fine but 48/32 combinations can yield 6); copying example code with packing hardcoded incorrectly.","solutions":["Use a power of 2: 1, 2, 4, or 8","If packing comes from bitwidth math, clamp/round down to the nearest power of two","For packing=0, default the parameter instead of passing it explicitly"],"exampleFix":"# before\nlayout = tcgen05.tmem_default_layout(packing=3)\n# after\nlayout = tcgen05.tmem_default_layout(packing=2)","handlingStrategy":"validation","validationCode":"assert isinstance(packing, int) and packing > 0 and packing.bit_count() == 1, packing","typeGuard":"def is_pow2(n: int) -> bool:\n    return isinstance(n, int) and n > 0 and n.bit_count() == 1","tryCatchPattern":null,"preventionTips":["Keep a single is_pow2 helper and use it wherever packing/bitwidth math appears","Round computed packing down to the nearest power of two: 1 << (x.bit_length() - 1)"],"tags":["jax","mosaic","tmem","layout","argument-validation"],"backgroundTag":"invalid-argument-value","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}