{"record":{"id":"34f7bd8f7b4fb98a","repo":"jax-ml/jax","slug":"expected-b-scales-to-have-a-m-128-layout-got-b-s","errorCode":null,"errorMessage":"Expected B scales to have a M=128 layout, got {b_scale.layout}","messagePattern":"Expected B scales to have a M=128 layout, got (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/tcgen05.py","lineNumber":481,"sourceCode":"          f\"Unsupported element type for block scaling: {a_element_type}\"\n      )\n    k_scales = k // scale_block\n    if a_scale.shape != (TMEM_ROWS, k_scales):\n      raise ValueError(\n          f\"A scale shape mismatch: expected ({TMEM_ROWS}, {k_scales}), got\"\n          f\" {a_scale.shape}\"\n      )\n    if a_scale.layout != scales_layout():\n      raise ValueError(f\"A scale layout {a_scale.layout} is not supported\")\n    if collective and m == 64:\n      if b_scale.layout != b_scales_m64_collective_layout():\n        raise ValueError(\n            \"Expected B scales to have a M=64 collective layout, got\"\n            f\" {b_scale.layout}\"\n        )\n    elif m == 128:\n      if b_scale.layout != scales_layout():\n        raise ValueError(\n            f\"Expected B scales to have a M=128 layout, got {b_scale.layout}\"\n        )\n    else:\n      raise AssertionError(\"Should not happen\")\n    if b_scale.shape[0] % 128 or b_scale.shape[0] < n * num_cta:\n      raise ValueError(\n          f\"B scale shape[0] must be a multiple of 128 and >= N={n * num_cta},\"\n          f\" got {b_scale.shape[0]}\"\n      )\n    if b_scale.shape[1] != k_scales:\n      raise ValueError(\n          f\"B scale shape mismatch: expected ({b_scale.shape[0]}, {k_scales}),\"\n          f\" got {b_scale.shape}\"\n      )\n  if is_sparse:\n    sparse_group_elems = 8 if utils.bitwidth(a_element_type) == 4 else 4\n    # Each sparse group has 2 entries.\n    expected_meta_k = k // sparse_group_elems * 2","sourceCodeStart":463,"sourceCodeEnd":499,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/tcgen05.py#L463-L499","documentation":"For block-scaled MMA with M=128, the B scale tensor must use the standard scales_layout(). This ValueError fires when any other layout is attached to b_scale in the m=128 path.","triggerScenarios":"Calling mma() with m=128 and b_scale.layout != scales_layout().","commonSituations":"Switching an M=64 collective kernel to M=128 without changing the B scale layout back; applying swizzled operand layouts to scales.","solutions":["Use scales_layout() for b_scale when m=128","If you need collective M=64, keep b_scales_m64_collective_layout() and m=64 consistent"],"exampleFix":"# before\nb_scale = TensorMemRefView(buf, shape, dt, layout=b_scales_m64_collective_layout())\n# after (m=128)\nb_scale = TensorMemRefView(buf, shape, dt, layout=scales_layout())","handlingStrategy":"validation","validationCode":"assert b_scale.layout == scales_layout() if m == 128 else True","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Keep B scale layout selection keyed on m and collective flags","Assert layout before entering the tile loop"],"tags":["gpu","mosaic","tcgen05","layout","block-scaling"],"backgroundTag":"invalid-memory-layout","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}