{"record":{"id":"0a2ef89f6434b072","repo":"jax-ml/jax","slug":"minor-dimension-of-shape-must-be-divisible-by-pack","errorCode":null,"errorMessage":"Minor dimension of shape must be divisible by packing, got: {shape}","messagePattern":"Minor dimension of shape must be divisible by packing, got: (.+?)","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/tcgen05.py","lineNumber":1069,"sourceCode":"        layout.warp_dims,\n        layout.lane_dims,\n        layout.vector_dim,\n        _check_canonical=False,\n    )\n\n  def as_tiled_layout(self) -> fa.TiledLayout:\n    return fa.TiledLayout(\n        self.tiling, self.warp_dims, self.lane_dims, self.vector_dim\n    )\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}\")","sourceCodeStart":1051,"sourceCodeEnd":1087,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/tcgen05.py#L1051-L1087","documentation":"Tensor Memory (TMEM) on Blackwell GPUs is addressed in packed columns; the inferred TMEM layout requires the minor (column) dimension of the 2D shape to be divisible by the packing factor. _infer_tmem_layout validates this before constructing a TMEMLayout. If shape[1] % packing != 0, no valid layout exists, so the error is raised.","triggerScenarios":"Calling infer_tmem_layout / from_alloc (with layout=None) or mma with a shape whose second dimension is not a multiple of the packing (e.g. packing=4 with shape=(128, 10)). Also triggered from _construct_smem_reftree or _default_tmem_layout_for_variable when a variable's column count is odd relative to packing.","commonSituations":"Mosaic GPU kernels using tcgen05 MMA where the N dimension of the accumulator is not padded to a power-of-2 multiple; using packing inferred from element bitwidth (e.g. packing=8 for i4/f8) with an N that isn't divisible by 8.","solutions":["Pad the minor dimension of the shape up to the next multiple of packing (e.g. N=10 with packing=4 -> N=12)","Reduce the packing factor so it divides the column count (packing=1 always works)","Pass an explicit layout= to from_alloc instead of relying on inference"],"exampleFix":"# before\nref = tcgen05.TMEMRef.from_alloc(alloc, (128, 10), collective=True, packing=4)\n# after\nref = tcgen05.TMEMRef.from_alloc(alloc, (128, 12), collective=True, packing=4)  # pad N to multiple of 4","handlingStrategy":"validation","validationCode":"def check_tmem_shape(shape, packing):\n    assert len(shape) == 2 and shape[1] % packing == 0, f'{shape} not divisible by packing={packing}'","typeGuard":"def has_valid_packing(shape: tuple[int, ...], packing: int) -> bool:\n    return len(shape) == 2 and shape[1] % packing == 0","tryCatchPattern":"try:\n    ref = tcgen05.TMEMRef.from_alloc(alloc, shape, collective=c)\nexcept ValueError as e:\n    if 'divisible by packing' in str(e):\n        shape = (shape[0], (shape[1] + packing - 1) // packing * packing)\n    else:\n        raise","preventionTips":["Derive packing from element bitwidth (32//bitwidth) so divisibility holds by construction","Pad N to a multiple of 8 at tile-size selection time","Centralize TMEM shape math in one helper with asserts"],"tags":["jax","mosaic","tmem","layout","shape-validation"],"backgroundTag":"shape-validation-failed","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}