{"record":{"id":"047c88adeafcdc61","repo":"jax-ml/jax","slug":"a-scale-shape-mismatch-expected-tmem-rows-k","errorCode":null,"errorMessage":"A scale shape mismatch: expected ({TMEM_ROWS}, {k_scales}), got {a_scale.shape}","messagePattern":"A scale shape mismatch: expected \\((.+?), (.+?)\\), got (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/tcgen05.py","lineNumber":467,"sourceCode":"      elif isinstance(scale_element_type, ir.Float8E8M0FNUType):\n        if base_scale_block not in (16, 32):\n          expected = \"32 or 64\" if is_sparse else \"16 or 32\"\n          raise ValueError(\n              f\"Scale block size mismatch: expected {expected}, got\"\n              f\" {scale_block}\"\n          )\n      else:\n        raise ValueError(\n            \"Scale element type mismatch: expected f8e8m0fnu or f8e4m3fn, got\"\n            f\" {scale_element_type}\"\n        )\n    else:\n      raise NotImplementedError(\n          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\")","sourceCodeStart":449,"sourceCodeEnd":485,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/tcgen05.py#L449-L485","documentation":"In block-scaled MMA, the A scale tensor must have exactly shape (TMEM_ROWS, k // scale_block). This error fires when a_scale.shape deviates, i.e. the scale tile doesn't cover K exactly once per scale block.","triggerScenarios":"Passing a_scale with shape != (TMEM_ROWS, k/scale_block), e.g. k=256, scale_block=32 but a_scale shape (128, 4) instead of (128, 8).","commonSituations":"Scale tensors computed for a different K dimension than the operands; off-by-one in k_scales computation in a code generator; padded K without padding the scales.","solutions":["Regenerate a_scale so its shape is (TMEM_ROWS, k // scale_block)","Ensure k is a multiple of scale_block and that the same k is used for operands and scales"],"exampleFix":"# before\nk, sb = 256, 32\na_scale = make_scale((TMEM_ROWS, 4))\n# after\na_scale = make_scale((TMEM_ROWS, k // sb))","handlingStrategy":"validation","validationCode":"TMEM_ROWS = 128\nassert a_scale.shape == (TMEM_ROWS, k // scale_block)","typeGuard":"def valid_a_scale(shape, k, scale_block, rows=128) -> bool:\n    return tuple(shape) == (rows, k // scale_block)","tryCatchPattern":null,"preventionTips":["Compute scale shapes from the same k/scale_block variables used for mma","Unit-test shape contracts for each kernel config"],"tags":["gpu","mosaic","tcgen05","shape-mismatch","block-scaling"],"backgroundTag":"tensor-shape-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}